Index

_

__init__() (nabla.core.array.Array method) (nabla.core.execution_context.ThreadSafeExecutionContext method) (nabla.nn.architectures.mlp.MLPBuilder method) (nabla.nn.architectures.MLPBuilder method) (nabla.nn.MLPBuilder method) (nabla.ops.operation.Operation method) (nabla.ops.operation.ReductionOperation method) (nabla.ops.operation.ViewOperation method) (nabla.transforms.utils.Trace method) (nabla.utils.types.Differentiable method)

A

abs() (in module nabla.ops.unary) accelerator() (in module nabla.utils.max_interop) accelerator_count() (in module nabla.utils.max_interop) accuracy() (in module nabla.nn) (in module nabla.nn.utils) (in module nabla.nn.utils.metrics) adam_step() (in module nabla.nn) (in module nabla.nn.optim) (in module nabla.nn.optim.adam) adamw_step() (in module nabla.nn) (in module nabla.nn.optim)

add() (in module nabla.ops.binary) add_arguments() (nabla.core.array.Array method) allclose() (in module nabla.utils.testing) arange() (in module nabla.ops.creation) argmax() (in module nabla.ops.reduce) args (nabla.core.array.Array attribute) Array (class in nabla.core.array) array() (in module nabla.ops.creation) array_slice() (in module nabla.ops.view) astype() (nabla.core.array.Array method) at() (nabla.core.array.Array method)

B

batch_dims (nabla.core.array.Array attribute) BINARY (nabla.utils.types.OperationType attribute) binary_cross_entropy_loss() (in module nabla.nn) (in module nabla.nn.losses) (in module nabla.nn.losses.classification)

BinaryOperation (class in nabla.ops.operation) broadcast_batch_dims() (in module nabla.ops.view) broadcast_to() (in module nabla.ops.view) build() (nabla.nn.architectures.mlp.MLPBuilder method) (nabla.nn.architectures.MLPBuilder method) (nabla.nn.MLPBuilder method)

C

cast() (in module nabla.ops.unary) clear() (nabla.core.execution_context.ThreadSafeExecutionContext method) compute_accuracy() (in module nabla.nn.utils) (in module nabla.nn.utils.training) compute_correlation() (in module nabla.nn.utils) (in module nabla.nn.utils.training) compute_node_hash() (in module nabla.core.graph_execution) (nabla.core.graph_execution.GraphTracer static method) compute_output_batch_dims() (nabla.ops.operation.BinaryOperation method) (nabla.ops.operation.ReductionOperation method) (nabla.ops.operation.UnaryOperation method) (nabla.ops.operation.ViewOperation method) compute_output_dtype() (nabla.ops.operation.BinaryOperation method) (nabla.ops.operation.UnaryOperation method) compute_output_shape() (nabla.ops.operation.BinaryOperation method) (nabla.ops.operation.Operation method) (nabla.ops.operation.ReductionOperation method) (nabla.ops.operation.UnaryOperation method) concatenate() (in module nabla.ops.view) cond() (in module nabla.ops.special) constant_schedule() (in module nabla.nn.optim) (in module nabla.nn.optim.schedules) contains() (nabla.core.execution_context.ThreadSafeExecutionContext method) copy_from() (nabla.core.array.Array method)

cos() (in module nabla.ops.unary) cosine_annealing_schedule() (in module nabla.nn) (in module nabla.nn.optim) (in module nabla.nn.optim.schedules) cotangent (nabla.core.array.Array attribute) cpu() (in module nabla.utils.max_interop) create_dataset() (in module nabla.nn) (in module nabla.nn.utils) (in module nabla.nn.utils.training) create_mlp_config() (in module nabla.nn) (in module nabla.nn.architectures) (in module nabla.nn.architectures.mlp) create_mlp_forward_and_loss() (in module nabla.nn.architectures) (in module nabla.nn.architectures.mlp) create_model() (nabla.core.graph_execution.ModelFactory static method) create_ones_like_cotangent() (in module nabla.utils.grad_utils) create_sin_dataset() (in module nabla.nn) (in module nabla.nn.utils) (in module nabla.nn.utils.training) CREATION (nabla.utils.types.OperationType attribute) cross_entropy_loss() (in module nabla.nn) (in module nabla.nn.losses) (in module nabla.nn.losses.classification) custom_kernel_path (nabla.core.array.Array attribute) custom_kernel_path() (nabla.ops.operation.Operation method)

D

decr_batch_dim_ctr() (in module nabla.ops.unary) device (nabla.core.array.Array attribute) device() (in module nabla.utils.max_interop) device_ref() (in module nabla.utils.max_interop) Differentiable (class in nabla.utils.types)

div() (in module nabla.ops.binary) djit() (in module nabla.transforms.jit) dropout() (in module nabla.nn) (in module nabla.nn.utils) (in module nabla.nn.utils.regularization) dtype (nabla.core.array.Array attribute)

E

EAGER (nabla.utils.types.ExecutionMode attribute) eagerxpr() (nabla.ops.operation.Operation method) elastic_net_regularization() (in module nabla.nn) (in module nabla.nn.utils) (in module nabla.nn.utils.regularization)

equal() (in module nabla.ops.binary) ExecutionMode (class in nabla.utils.types) exp() (in module nabla.ops.unary) exponential_decay_schedule() (in module nabla.nn) (in module nabla.nn.optim) (in module nabla.nn.optim.schedules)

F

f1_score() (in module nabla.nn) (in module nabla.nn.utils) (in module nabla.nn.utils.metrics) floor() (in module nabla.ops.unary) floordiv() (in module nabla.ops.binary) format_dtype() (in module nabla.utils.formatting)

format_shape_and_dtype() (in module nabla.utils.formatting) forward() (nabla.ops.operation.BinaryOperation method) (nabla.ops.operation.Operation method) (nabla.ops.operation.UnaryOperation method) from_impl() (nabla.core.array.Array class method) from_numpy() (nabla.core.array.Array class method) full_like() (in module nabla.ops.creation)

G

gelu() (in module nabla.nn) (in module nabla.nn.layers) (in module nabla.nn.layers.activations) get() (nabla.core.execution_context.ThreadSafeExecutionContext method) get_activation() (in module nabla.nn.layers) (in module nabla.nn.layers.activations) get_arguments() (nabla.core.array.Array method) get_broadcasted_shape() (in module nabla.utils.shape_utils) get_or_create() (nabla.core.execution_context.ThreadSafeExecutionContext method)

get_trace() (in module nabla.core.graph_execution) (nabla.core.graph_execution.GraphTracer static method) get_traced_nodes() (nabla.transforms.utils.Trace method) glorot_uniform() (in module nabla.ops.creation) grad() (in module nabla.transforms.grad) gradient_clipping() (in module nabla.nn) (in module nabla.nn.utils) (in module nabla.nn.utils.regularization) GraphTracer (class in nabla.core.graph_execution) greater_equal() (in module nabla.ops.binary)

H

he_normal() (in module nabla.nn) (in module nabla.nn.init) (in module nabla.nn.init.variance_scaling) (in module nabla.ops.creation) he_uniform() (in module nabla.nn.init) (in module nabla.nn.init.variance_scaling) (in module nabla.ops.creation)

huber_loss() (in module nabla.nn) (in module nabla.nn.losses) (in module nabla.nn.losses.regression)

I

impl (nabla.core.array.Array property) impl_() (nabla.core.array.Array method) incr_batch_dim_ctr() (in module nabla.ops.unary) init_adam_state() (in module nabla.nn) (in module nabla.nn.optim) (in module nabla.nn.optim.adam) init_adamw_state() (in module nabla.nn) (in module nabla.nn.optim) (in module nabla.nn.optim.adamw)

init_sgd_state() (in module nabla.nn) (in module nabla.nn.optim) (in module nabla.nn.optim.sgd) initialize_mlp_params() (in module nabla.nn) (in module nabla.nn.init) (in module nabla.nn.init.variance_scaling) is_documentable() (in module nabla.utils) (in module nabla.utils.docs)

J

jacfwd() (in module nabla.transforms.jacfwd) jacrev() (in module nabla.transforms.jacrev) jit() (in module nabla.transforms.jit)

jvp() (in module nabla.transforms.jvp) jvp_rule (nabla.core.array.Array attribute) jvp_rule() (nabla.ops.operation.Operation method) (nabla.utils.types.Differentiable method)

K

kernel_impl_path (nabla.core.array.Array attribute)

L

l1_regularization() (in module nabla.nn) (in module nabla.nn.utils) (in module nabla.nn.utils.regularization) l2_regularization() (in module nabla.nn) (in module nabla.nn.utils) (in module nabla.nn.utils.regularization) LAZY (nabla.utils.types.ExecutionMode attribute) leaky_relu() (in module nabla.nn) (in module nabla.nn.layers) (in module nabla.nn.layers.activations) learning_rate_schedule() (in module nabla.nn) (in module nabla.nn.optim) (in module nabla.nn.optim.schedules)

lecun_normal() (in module nabla.nn) (in module nabla.nn.init) (in module nabla.nn.init.variance_scaling) (in module nabla.ops.creation) lecun_uniform() (in module nabla.ops.creation) linear_forward() (in module nabla.nn) (in module nabla.nn.layers) (in module nabla.nn.layers.linear) log() (in module nabla.ops.unary) log_softmax() (in module nabla.nn) (in module nabla.nn.layers) (in module nabla.nn.layers.activations) logical_not() (in module nabla.ops.unary) logsumexp() (in module nabla.ops.special)

M

make_staged_pytree() (in module nabla.transforms.utils) make_traced_pytree() (in module nabla.transforms.utils) make_unstaged_pytree() (in module nabla.transforms.utils) make_untraced_pytree() (in module nabla.transforms.utils) matmul() (in module nabla.ops.linalg) max() (in module nabla.ops.reduce) maximum() (in module nabla.ops.binary) maxpr (nabla.core.array.Array attribute) maxpr() (nabla.ops.operation.Operation method) mean() (in module nabla.ops.reduce) mean_absolute_error() (in module nabla.nn) (in module nabla.nn.losses) (in module nabla.nn.losses.regression) mean_absolute_error_metric() (in module nabla.nn) (in module nabla.nn.utils) (in module nabla.nn.utils.metrics) mean_squared_error() (in module nabla.nn) (in module nabla.nn.losses) (in module nabla.nn.losses.regression) mean_squared_error_metric() (in module nabla.nn) (in module nabla.nn.utils) (in module nabla.nn.utils.metrics) minimum() (in module nabla.ops.binary) mlp_forward() (in module nabla.nn) (in module nabla.nn.layers) (in module nabla.nn.layers.linear) mlp_forward_with_activations() (in module nabla.nn) (in module nabla.nn.layers) (in module nabla.nn.layers.linear) MLPBuilder (class in nabla.nn) (class in nabla.nn.architectures) (class in nabla.nn.architectures.mlp) mod() (in module nabla.ops.binary) ModelFactory (class in nabla.core.graph_execution) module nabla.core nabla.core.array nabla.core.execution_context nabla.core.graph_execution nabla.nn nabla.nn.architectures nabla.nn.architectures.mlp nabla.nn.init nabla.nn.init.variance_scaling nabla.nn.layers nabla.nn.layers.activations nabla.nn.layers.linear nabla.nn.losses nabla.nn.losses.classification nabla.nn.losses.regression nabla.nn.optim nabla.nn.optim.adam nabla.nn.optim.adamw nabla.nn.optim.schedules nabla.nn.optim.sgd nabla.nn.utils nabla.nn.utils.metrics nabla.nn.utils.regularization nabla.nn.utils.training nabla.ops nabla.ops.binary nabla.ops.creation nabla.ops.kernels nabla.ops.linalg nabla.ops.operation nabla.ops.reduce nabla.ops.special nabla.ops.unary nabla.ops.view nabla.transforms nabla.transforms.grad nabla.transforms.jacfwd nabla.transforms.jacrev nabla.transforms.jit nabla.transforms.jvp nabla.transforms.utils nabla.transforms.vjp nabla.transforms.vmap nabla.utils nabla.utils.docs nabla.utils.formatting nabla.utils.grad_utils nabla.utils.max_interop nabla.utils.shape_utils nabla.utils.testing nabla.utils.types

move_axis_from_front() (in module nabla.ops.view) move_axis_from_front_of_batch_dims() (in module nabla.ops.view) move_axis_to_front() (in module nabla.ops.view) move_axis_to_front_of_batch_dims() (in module nabla.ops.view) move_to_best_device() (in module nabla.ops.operation) mul() (in module nabla.ops.binary)

N

nabla.core module nabla.core.array module nabla.core.execution_context module nabla.core.graph_execution module nabla.nn module nabla.nn.architectures module nabla.nn.architectures.mlp module nabla.nn.init module nabla.nn.init.variance_scaling module nabla.nn.layers module nabla.nn.layers.activations module nabla.nn.layers.linear module nabla.nn.losses module nabla.nn.losses.classification module nabla.nn.losses.regression module nabla.nn.optim module nabla.nn.optim.adam module nabla.nn.optim.adamw module nabla.nn.optim.schedules module nabla.nn.optim.sgd module nabla.nn.utils module nabla.nn.utils.metrics module nabla.nn.utils.regularization module nabla.nn.utils.training module nabla.ops module nabla.ops.binary module nabla.ops.creation module nabla.ops.kernels module

nabla.ops.linalg module nabla.ops.operation module nabla.ops.reduce module nabla.ops.special module nabla.ops.unary module nabla.ops.view module nabla.transforms module nabla.transforms.grad module nabla.transforms.jacfwd module nabla.transforms.jacrev module nabla.transforms.jit module nabla.transforms.jvp module nabla.transforms.utils module nabla.transforms.vjp module nabla.transforms.vmap module nabla.utils module nabla.utils.docs module nabla.utils.formatting module nabla.utils.grad_utils module nabla.utils.max_interop module nabla.utils.shape_utils module nabla.utils.testing module nabla.utils.types module name (nabla.core.array.Array attribute) ndarange() (in module nabla.ops.creation) ndarange_like() (in module nabla.ops.creation) negate() (in module nabla.ops.unary) no_doc() (in module nabla.utils.docs) nodoc() (in module nabla.utils) (in module nabla.utils.docs) not_equal() (in module nabla.ops.binary)

O

ones() (in module nabla.ops.creation) ones_like() (in module nabla.ops.creation)

Operation (class in nabla.ops.operation) OperationType (class in nabla.utils.types)

P

pad() (in module nabla.ops.view) pearson_correlation() (in module nabla.nn) (in module nabla.nn.utils) (in module nabla.nn.utils.metrics) permute() (in module nabla.ops.view) (nabla.core.array.Array method)

permute_batch_dims() (in module nabla.ops.view) pow() (in module nabla.ops.binary) precision() (in module nabla.nn) (in module nabla.nn.utils) (in module nabla.nn.utils.metrics) pullback() (in module nabla.transforms.utils) pushfwd() (in module nabla.transforms.utils)

R

r_squared() (in module nabla.nn) (in module nabla.nn.utils) (in module nabla.nn.utils.metrics) rand() (in module nabla.ops.creation) rand_like() (in module nabla.ops.creation) randn() (in module nabla.ops.creation) randn_like() (in module nabla.ops.creation) realize() (nabla.core.array.Array method) realize_() (in module nabla.core.graph_execution) recall() (in module nabla.nn) (in module nabla.nn.utils) (in module nabla.nn.utils.metrics)

REDUCTION (nabla.utils.types.OperationType attribute) ReductionOperation (class in nabla.ops.operation) relu() (in module nabla.nn) (in module nabla.nn.layers) (in module nabla.nn.layers.activations) (in module nabla.ops.unary) reshape() (in module nabla.ops.view) (nabla.core.array.Array method)

S

select_gradients_by_argnums() (in module nabla.utils.grad_utils) set() (nabla.core.array.Array method) (nabla.core.execution_context.ThreadSafeExecutionContext method) set_maxpr() (nabla.core.array.Array method) sgd_step() (in module nabla.nn) (in module nabla.nn.optim) (in module nabla.nn.optim.sgd) shallow_copy() (in module nabla.ops.view) shape (nabla.core.array.Array attribute) should_document() (in module nabla.utils) (in module nabla.utils.docs) sigmoid() (in module nabla.nn) (in module nabla.nn.layers) (in module nabla.nn.layers.activations) (in module nabla.ops.unary) silu() (in module nabla.nn) (in module nabla.nn.layers) (in module nabla.nn.layers.activations) sin() (in module nabla.ops.unary) size (nabla.core.array.Array property) size() (nabla.core.execution_context.ThreadSafeExecutionContext method) skip_doc() (in module nabla.utils.docs) softmax() (in module nabla.nn) (in module nabla.nn.layers) (in module nabla.nn.layers.activations) (in module nabla.ops.special)

softmax_cross_entropy_loss() (in module nabla.nn) (in module nabla.nn.losses) (in module nabla.nn.losses.classification) sparse_cross_entropy_loss() (in module nabla.nn) (in module nabla.nn.losses) (in module nabla.nn.losses.classification) spectral_normalization() (in module nabla.nn.utils) (in module nabla.nn.utils.regularization) sqrt() (in module nabla.ops.unary) squeeze() (in module nabla.ops.view) squeeze_batch_dims() (in module nabla.ops.view) stack() (in module nabla.ops.view) stage_realization (nabla.core.array.Array attribute) step_decay_schedule() (in module nabla.nn.optim) (in module nabla.nn.optim.schedules) sub() (in module nabla.ops.binary) sum() (in module nabla.ops.reduce) (nabla.core.array.Array method) sum_batch_dims() (in module nabla.ops.reduce) swish() (in module nabla.nn) (in module nabla.nn.layers) (in module nabla.nn.layers.activations)

T

tangent (nabla.core.array.Array attribute) tanh() (in module nabla.nn) (in module nabla.nn.layers) (in module nabla.nn.layers.activations) (in module nabla.ops.unary) tensor_value (nabla.core.array.Array attribute) ThreadSafeExecutionContext (class in nabla.core.execution_context) to() (nabla.core.array.Array method) to_numpy() (nabla.core.array.Array method) top_k_accuracy() (in module nabla.nn.utils) (in module nabla.nn.utils.metrics)

Trace (class in nabla.transforms.utils) trace_function() (nabla.transforms.utils.Trace class method) traced (nabla.core.array.Array attribute) transfer_to() (in module nabla.ops.unary) transpose() (in module nabla.ops.view) (nabla.core.array.Array method) tree_flatten() (in module nabla.transforms.utils) tree_map() (in module nabla.transforms.utils) tree_unflatten() (in module nabla.transforms.utils) triu() (in module nabla.ops.creation)

U

UNARY (nabla.utils.types.OperationType attribute) UnaryOperation (class in nabla.ops.operation)

unsqueeze() (in module nabla.ops.view) unsqueeze_batch_dims() (in module nabla.ops.view)

V

validate_scalar_output() (in module nabla.utils.grad_utils) value_and_grad() (in module nabla.transforms.grad) VIEW (nabla.utils.types.OperationType attribute) ViewOperation (class in nabla.ops.operation) visited (nabla.core.array.Array attribute)

vjp() (in module nabla.transforms.vjp) vjp_rule (nabla.core.array.Array attribute) vjp_rule() (nabla.ops.operation.Operation method) (nabla.utils.types.Differentiable method) vmap() (in module nabla.transforms.vmap)

W

warmup_cosine_schedule() (in module nabla.nn) (in module nabla.nn.optim) (in module nabla.nn.optim.schedules) where() (in module nabla.ops.special) with_activation() (nabla.nn.architectures.mlp.MLPBuilder method) (nabla.nn.architectures.MLPBuilder method) (nabla.nn.MLPBuilder method) with_final_activation() (nabla.nn.architectures.mlp.MLPBuilder method) (nabla.nn.architectures.MLPBuilder method) (nabla.nn.MLPBuilder method)

with_init_method() (nabla.nn.architectures.mlp.MLPBuilder method) (nabla.nn.architectures.MLPBuilder method) (nabla.nn.MLPBuilder method) with_layers() (nabla.nn.architectures.mlp.MLPBuilder method) (nabla.nn.architectures.MLPBuilder method) (nabla.nn.MLPBuilder method) with_seed() (nabla.nn.architectures.mlp.MLPBuilder method) (nabla.nn.architectures.MLPBuilder method) (nabla.nn.MLPBuilder method)

X

xavier_normal() (in module nabla.nn) (in module nabla.nn.init) (in module nabla.nn.init.variance_scaling) (in module nabla.ops.creation)

xavier_uniform() (in module nabla.nn.init) (in module nabla.nn.init.variance_scaling) (in module nabla.ops.creation) xpr() (in module nabla.transforms.utils)

Z

zeros() (in module nabla.ops.creation)

zeros_like() (in module nabla.ops.creation)