List of all items
Structs
- Array
- At
- Complex64
- CustomFunction
- CustomFunctionFlat
- Device
- Ellipsis
- GatherQmm
- MappedFile
- NewAxis
- Quant
- Quantized
- ScanOptions
- Scheduled
- Stepped
- Stream
- VarOptions
- Vmap
- bf16
- distributed::Group
- f16
- fast::AttentionOptions
- fast::AttentionPlan
- fast::KernelCall
- fast::KernelConfig
- fast::MetalKernel
- fast::Rotation
- fast::StepOptions
- graph::CustomCall
- graph::Flags
- graph::Fused
- graph::SliceArg
- tree::Param
- tree::TreeDef
- tree::Visit
- tree::VisitMut
Enums
- Category
- DType
- DeviceKind
- Error
- OpErrorKind
- QuantMode
- Scalar
- ScalarOrArray
- distributed::AllReduce
- distributed::CollectiveOp
- fast::AttentionKernel
- fast::AttentionMask
- fast::FastOp
- fast::TemplateArg
- graph::ArgReduceOp
- graph::BinaryOp
- graph::FusedRef
- graph::Primitive
- graph::ReduceOp
- graph::ScanOp
- graph::ScatterReduce
- graph::UnaryOp
- tree::LeafKind
- tree::LoadMode
Traits
- ArrayIndex
- Axes
- Axis
- Element
- GradOutput
- InAxes
- InAxis
- IntoResult
- IntoShape
- OutAxes
- ScalarOperand
- Step
- distributed::Communicator
- tree::BufferLeaf
- tree::Leaves
- tree::Static
- tree::Tree
- tree::Visitor
- tree::VisitorMut
Macros
Functions
- async_eval
- at_least_float
- checkpoint
- checkpoint_flat
- custom_function
- custom_function_flat
- default_device
- default_stream
- distributed::all_gather
- distributed::all_max
- distributed::all_min
- distributed::all_sum
- distributed::default_group
- distributed::recv
- distributed::recv_like
- distributed::send
- distributed::set_default_group
- distributed::sum_gradients
- eval
- fast::add_rms_norm
- fast::cross_entropy
- fast::gated_delta
- fast::gated_delta_slots
- fast::gated_delta_steps
- fast::gated_delta_steps_with
- fast::layer_norm
- fast::norm_rope
- fast::norm_rope_kv
- fast::paged_attention
- fast::paged_write
- fast::plan_attention
- fast::rms_norm
- fast::rope
- fast::scaled_dot_product_attention
- grad
- grad_flat
- graph::broadcast_shapes
- graph::flags
- graph::from_primitive
- graph::id
- graph::inputs
- graph::is_tracer
- graph::output_index
- graph::placeholder
- graph::primitive
- graph::siblings
- graph::slice
- graph::strides
- jvp
- jvp_flat
- ops::abs
- ops::add
- ops::addmm
- ops::all
- ops::all_keepdims
- ops::any
- ops::any_keepdims
- ops::arange
- ops::argmax
- ops::argmax_keepdims
- ops::argmin
- ops::argmin_keepdims
- ops::argpartition
- ops::argsort
- ops::astype
- ops::broadcast_to
- ops::concatenate
- ops::conjugate
- ops::contiguous
- ops::conv1d
- ops::copy
- ops::cos
- ops::cummax
- ops::cummax_with
- ops::cummin
- ops::cummin_with
- ops::cumprod
- ops::cumprod_with
- ops::cumsum
- ops::cumsum_with
- ops::dequantize
- ops::divide
- ops::equal
- ops::erfinv
- ops::exp
- ops::expand_dims
- ops::floor_divide
- ops::from_fp8
- ops::full
- ops::full_like
- ops::gather_mm
- ops::gather_qmm
- ops::greater
- ops::greater_equal
- ops::less
- ops::less_equal
- ops::log
- ops::logaddexp
- ops::logcumsumexp
- ops::logcumsumexp_with
- ops::logsumexp
- ops::logsumexp_keepdims
- ops::matmul
- ops::max
- ops::max_keepdims
- ops::maximum
- ops::mean
- ops::mean_keepdims
- ops::min
- ops::min_keepdims
- ops::minimum
- ops::multiply
- ops::negative
- ops::not_equal
- ops::ones
- ops::ones_like
- ops::partition
- ops::power
- ops::prod
- ops::prod_keepdims
- ops::put_along_axis
- ops::quantize
- ops::quantized_matmul
- ops::repeat
- ops::reshape
- ops::scatter_add_along_axis
- ops::sigmoid
- ops::sign
- ops::sin
- ops::slice_update
- ops::softmax
- ops::sort
- ops::split
- ops::split_at
- ops::split_n
- ops::sqrt
- ops::squeeze
- ops::stack
- ops::std
- ops::std_keepdims
- ops::std_with
- ops::stop_gradient
- ops::subtract
- ops::sum
- ops::sum_keepdims
- ops::take
- ops::take_along_axis
- ops::tanh
- ops::to_fp8
- ops::topk
- ops::transpose
- ops::var
- ops::var_keepdims
- ops::var_with
- ops::view
- ops::where_
- ops::zeros
- ops::zeros_like
- promote_types
- random::bernoulli
- random::bits
- random::categorical
- random::gumbel
- random::key
- random::next_key
- random::normal
- random::randint
- random::seed
- random::split
- random::split_n
- random::uniform
- row_exact
- schedule
- set_default_stream
- speculative::accept_greedy
- speculative::accept_sampled
- speculative::read
- tree::collect_leaves
- value_and_grad
- value_and_grad_flat
- vjp
- vjp_flat
- vmap
- vmap_flat
- with_row_exact
- with_stream