Models Agents Infrastructure Pricing Docs Blog Security

Models Agents Infrastructure Pricing Docs Blog Security

Thunder

Platform

Lightning AIDevelop, train and deploy on Lightning AI Python SDKProgrammatic access to the Lightning AI platform CLIAccess the Lightning AI platform from the command line

Open Source Libraries

PyTorch LightningFinetune and pretrain AI models. Focus on science, not engineering.Lightning FabricScale foundation models to 1000s of GPUs with expert-level control TorchMetrics90+ PyTorch metrics, optimized for distributed training ThunderCompiler for PyTorch: Make PyTorch models faster!LitLoggerExperiment tracking built into Lightning AI LitDataStream and process datasets at scale LitServeEasily serve AI models — Lightning fast LitAILLM router and minimal agent framework in one LitGPT20+ high-performance LLMs with recipes to pretrain, finetune and deploy at scale.

Stable (0.2.6)

Home

Install Hello World Using examine

Basic

Overview Thunder step by step The sharp edges Train a MLP on MNIST FAQ

Intermediate

Additional executors Distributed Data Parallel What's next Benchmarking Thunder

Advanced

Inside Thunder Extending Thunder Contributing to Thunder

API reference

Baseutils Codeutils Devices Dtypes Langctxs Prims Proxies PyTree Rematerialization Symbol Trace Transforms

Clang Examine Distributed

APEX Executor cuDNN Executor nvFuser Executor Optimization Passes Torch Compile Executor Triton Executor Utils

Torch Extend Transforms Dynamo Recipes Plugins

thunder.cache_hits thunder.cache_misses thunder.cache_option thunder.clang.abs thunder.clang.acos thunder.clang.acosh thunder.clang.add thunder.clang.arange thunder.clang.asin thunder.clang.asinh thunder.clang.atan thunder.clang.atan2 thunder.clang.atanh thunder.clang.bitwise_and thunder.clang.bitwise_not thunder.clang.bitwise_or thunder.clang.bitwise_xor thunder.clang.cat thunder.clang.ceil thunder.clang.compute_broadcast_shape thunder.clang.convolution thunder.clang.copysign thunder.clang.cos thunder.clang.cosh thunder.clang.device_put thunder.clang.diagonal thunder.clang.eq thunder.clang.erf thunder.clang.erfc thunder.clang.erfcinv thunder.clang.erfinv thunder.clang.exp thunder.clang.exp2 thunder.clang.expand thunder.clang.expm1 thunder.clang.flatten thunder.clang.floor thunder.clang.floor_divide thunder.clang.fmod thunder.clang.full thunder.clang.full_like thunder.clang.ge thunder.clang.gt thunder.clang.index_add thunder.clang.isfinite thunder.clang.le thunder.clang.lgamma thunder.clang.log thunder.clang.log1p thunder.clang.log2 thunder.clang.log10 thunder.clang.logical_and thunder.clang.lt thunder.clang.matrix_transpose thunder.clang.maybe_broadcast thunder.clang.maybe_convert_to_dtype thunder.clang.mod thunder.clang.movedim thunder.clang.mul thunder.clang.ndtri thunder.clang.ne thunder.clang.neg thunder.clang.nextafter thunder.clang.pow thunder.clang.reciprocal thunder.clang.remainder thunder.clang.reshape thunder.clang.round thunder.clang.rsqrt thunder.clang.scatter_add thunder.clang.sigmoid thunder.clang.sign thunder.clang.signbit thunder.clang.silu thunder.clang.sin thunder.clang.sinh thunder.clang.slice_in_dim thunder.clang.sqrt thunder.clang.squeeze thunder.clang.stack thunder.clang.stride_order thunder.clang.sub thunder.clang.take thunder.clang.take_along_axis thunder.clang.tan thunder.clang.tanh thunder.clang.transpose thunder.clang.true_divide thunder.clang.trunc thunder.clang.uniform thunder.clang.uniform_like thunder.clang.unsqueeze thunder.clang.where thunder.common.CACHE_OPTIONS thunder.common.CompileData thunder.common.CompileStats thunder.compile_data thunder.compile_stats thunder.core.baseutils.BoundSymbolInterface thunder.core.baseutils.check thunder.core.baseutils.check_type thunder.core.baseutils.NumberProxyInterface thunder.core.baseutils.ProxyInterface thunder.core.baseutils.SymbolInterface thunder.core.baseutils.TensorProxyInterface thunder.core.codeutils.get_siginfo thunder.core.codeutils.prettyprint thunder.core.codeutils.SigInfo thunder.core.devices.Device thunder.core.devices.device_from_string thunder.core.devices.DeviceType thunder.core.devices.devicetype_string thunder.core.devices.to_device thunder.core.dtypes.bfloat16 thunder.core.dtypes.bool8 thunder.core.dtypes.complex32 thunder.core.dtypes.complex64 thunder.core.dtypes.complex128 thunder.core.dtypes.dtype thunder.core.dtypes.float16 thunder.core.dtypes.float32 thunder.core.dtypes.float64 thunder.core.dtypes.int8 thunder.core.dtypes.int16 thunder.core.dtypes.int32 thunder.core.dtypes.int64 thunder.core.dtypes.to_dtype thunder.core.dtypes.uint8 thunder.core.langctxs.get_langctx thunder.core.langctxs.langctx thunder.core.langctxs.reset_langctx thunder.core.langctxs.set_langctx thunder.core.prims.PrimIDs thunder.core.proxies.CollectionProxy thunder.core.proxies.ComplexProxy thunder.core.proxies.FloatProxy thunder.core.proxies.FutureTensorProxy thunder.core.proxies.IntegerProxy thunder.core.proxies.is_proxyable thunder.core.proxies.proxy thunder.core.proxies.pytype thunder.core.proxies.pyval thunder.core.proxies.TensorProxy thunder.core.proxies.unvariableify thunder.core.proxies.Variable thunder.core.proxies.variableify thunder.core.rematerialization.rematerialize thunder.core.symbol.BoundSymbol thunder.core.symbol.BoundSymbolRHS thunder.core.symbol.Symbol thunder.core.trace.from_trace thunder.core.trace.get_tracectx thunder.core.trace.reset_tracectx thunder.core.trace.set_tracectx thunder.core.trace.TraceCtx thunder.core.transforms.bsym_list_to_dag thunder.core.transforms.insert_inplace thunder.core.transforms.Node thunder.core.transforms.replace_inplace thunder.core.transforms.toposort_bsym_dag thunder.core.transforms.VISIT_TYPE thunder.core.transforms.visitor_transform thunder.DebugOptions thunder.distributed.column_parallel thunder.distributed.ddp thunder.distributed.fsdp thunder.distributed.FSDPBucketingStrategy thunder.distributed.FSDPType thunder.distributed.get_skip_data_parallel_grad_sync thunder.distributed.reset_skip_data_parallel_grad_sync thunder.distributed.row_parallel thunder.distributed.set_skip_data_parallel_grad_sync thunder.distributed.skip_data_parallel_grad_sync thunder.examine.examine thunder.executors.apexex.apex_entropy_available thunder.executors.cudnnex.cudnn_ex thunder.executors.get_nvfuser_executor thunder.executors.get_torch_executor thunder.executors.nvfuser_available thunder.executors.passes.del_last_used thunder.executors.passes.transform_for_execution thunder.executors.passes.update_fusion_call_ctx thunder.executors.torch_compile.make_compiled thunder.executors.torch_compile.TorchCompileExecutor thunder.executors.utils.Region thunder.jit thunder.last_backward_traces thunder.last_compile_options thunder.last_interpreted_instructions thunder.last_interpreter_log thunder.last_prologue_traces thunder.last_traces thunder.list_transforms thunder.plugins.DDP thunder.plugins.FP8 thunder.plugins.FSDP thunder.plugins.QuantizeInt4 thunder.plugins.ReduceOverhead thunder.recipes.BaseRecipe thunder.recipes.HFTransformers thunder.torch._register_custom_op thunder.torch.size thunder.torch.to_float thunder.torch.torchsymbol thunder.transforms.ConstantFolding thunder.transforms.MaterializationTransform

thunder.dynamo.ThunderCompiler thunder.dynamo.thunderfx thunder.dynamo.ThunderFXCompiledObject thunder.extend.add_always_executor thunder.extend.add_default_executor thunder.extend.deregister_executor thunder.extend.Executor thunder.extend.get_all_executors thunder.extend.get_always_executors thunder.extend.get_default_executors thunder.extend.get_executor thunder.extend.register_executor thunder.extend.remove_always_executor thunder.extend.remove_default_executor thunder.extend.set_always_executors thunder.extend.set_default_executors

Docs/Stable (0.2.6)/API reference/Generated/thunder.clang.take_along_axis

thunder.clang.take_along_axis

thunder.clang.take_along_axis(a, /, indices, dim)[source]

Return type:

Parameters:

a (TensorProxy) –
indices (TensorProxy) –
dim (int) –

Was this page helpful?

Thanks for the feedback.

What went wrong?

Inaccurate or out of dateHard to understandMissing informationHard to find things / poorly organizedCode sample didn't workSomething is broken (link, image, layout)

Share with AI

Open in ChatGPT Open in Claude

Was this helpful?

Thanks for the feedback.

What went wrong?

Inaccurate or out of dateHard to understandMissing informationHard to find things / poorly organizedCode sample didn't workSomething is broken (link, image, layout)

↑↓ Navigate↵ Open⌘↵ New tabesc Close