Operation Reference =================== .. toctree:: :maxdepth: 1 :caption: Operations A Abs Acos Acosh ActivityRegularization Adadelta Adafactor Adagrad Adam Adamax Adamw AdaptiveAveragePooling1D AdaptiveAveragePooling2D AdaptiveAveragePooling3D AdaptiveGradClipState AdaptiveMaxPooling1D AdaptiveMaxPooling2D AdaptiveMaxPooling3D Add AddDecayedWeightsState AddNoiseState AdditiveAttention All AllClose Alphadropout Angle Any Append ApplyAlongAxis ApplyIfFiniteState Arange ArgMax ArgMin ArgPartition ArgSort ArgWhere Array ArrayEqual ArrayEquiv ArrayLike ArraySplit AsArray Asin Asinh AssertClose Astype Atan Atan2 Atanh Atleast1d Atleast2d Atleast3d Atom Attention AugMix AutoContrast AuxOutput Average AveragePooling1D AveragePooling2D AveragePooling3D AvgPool1d AvgPool2d AvgPool3d Avgpool B BLOCK_M Bartlett BatchNorm BatchNormalization Bernoulli Bfloat16 Bidirectional Bilinear Bincount BitwiseAnd BitwiseCount BitwiseInvert BitwiseLeftShift BitwiseNot BitwiseOr BitwiseRightShift BitwiseXor Blackman Block Bool BroadcastArrays BroadcastShapes BroadcastTo C CTCLoss Cache Call CanCast Cast CastBool CastByte CastChar CastDouble CastFloat CastHalf CastInt CastLong CastShort CategoryEncoding Cbrt Ceil Celu CenterCrop Choose Clamp Clip ClipByGlobalNormState ClipGradNorm ClipState ColumnStack Compile Complex128 Complex64 ComplexWarning Compress ComputeCv Concat Concatenate ConditionallyMaskState ConditionallyTransformState Conj Conjugate Constant ControlVariate Conv Conv1DTranspose Conv1d Conv2DTranspose Conv2d Conv3DTranspose Conv3d ConvLSTM1D ConvLSTM2D ConvLSTM3D ConvT ConvTranspose1d ConvTranspose2d ConvTranspose3d Convolution1D Convolution1DTranspose Convolution2D Convolution2DTranspose Convolution3D Convolution3DTranspose Convolve Copy Copysign Corrcoef Correlate Cos Cosh CosineAnnealingLR CosineSimilarity CountNonzero Cov Cropping1D Cropping2D Cropping3D Cross CrossEntropyLoss Csingle CudaAvailable CumSum Cumprod CumulativeProd CumulativeSum CutMix CvExpectedValue CvState DataLoader DataParallel Deg2rad Degrees Delete Dense DepthwiseConv1D DepthwiseConv2D Det Device Diag DiagIndices DiagIndicesFrom Diagflat Diagonal Diff Digitize Discretization Div Divide Divmod Dot Double Dropout Dropout2d Dropout3d Dsplit Dstack Dtype E Ediff1d Einsum EinsumDense EinsumPath Elu EmaState Embedding Empty EmptyLike Equal Equalization Erf EulerGamma Exp Exp2 ExpandDims Expm1 Extract Eye Fabs FillDiagonal Finfo Fix Flatnonsense Flatnonzero Flatten FlaxLayer Flexible Flip Fliplr Flipud Float16 Float32 Float4_e2m1fn Float64 Float8_e3m4 Float8_e4m3 Float8_e4m3b11fnuz Float8_e4m3fn Float8_e4m3fnuz Float8_e5m2 Float8_e5m2fnuz Float8_e8m0fnu FloatPower Float_ Floating Floor FloorDivide Fmax Fmin Fmod Frexp FromBuffer FromDlpack FromFile FromFunction FromIter FromPyFunc FromString Full FullLike FusedMultiplyAdd GELU GRU Gather GaussianDropout GaussianNoise Gaussiannll Gcd Generic Geomspace GetPrintoptions GlobalAveragePooling1D GlobalAveragePooling2D GlobalAveragePooling3D GlobalAvgPool1D GlobalAvgPool2D GlobalAvgPool3D GlobalMaxPool1D GlobalMaxPool2D GlobalMaxPool3D GlobalMaxPooling1D GlobalMaxPooling2D GlobalMaxPooling3D Glu Gradient Grayscale Greater GreaterEqual Group GroupNormalization GroupQueryAttention Groupnorm Grucell Hamming Hanning HardSigmoid HardSwish Hardshrink Hardtanh HashedCrossing Hashing Heaviside Histogram Histogram2d HistogramBinEdges Histogramdd Hsplit Hstack Hypot I0 INTEGER Identity Imag Indices InjectHyperparamsState InjectStatefulHyperparamsState Inner InputLayer InputSpec Insert Instancenorm Int16 Int32 Int64 Int8 IntegerLookup Interp Intersect1d Invert IsComplex IsComplexObj IsDtype IsFinite IsIn IsInf IsNaN IsNegInf IsPosInf IsReal IsRealObj IsScalar IsSubDtype Isclose Iterable Ix_ JaxLayer Kaiser Kldiv Kron L1 LOWER_RIGHT LSTM Lambda Layer LayerNorm LayerNormalization LazyConvTranspose1d LazyConvTranspose2d LazyConvTranspose3d Lcm Ldexp LeakyRelu LeftShift Less LessEqual Lexsort Linear LinearT Linspace Lion List Load Log Log10 Log1p Log2 LogSoftmax Logaddexp Logaddexp2 Logcosh LogicalAnd LogicalNot LogicalOr LogicalXor Logsigmoid Logspace Logsumexp Lstmcell M MA MSELoss MarginRankingLoss Marginranking MaskIndices MaskedNode MaskedState Masking MatMul MatrixTranspose Matvec Max MaxNumBoundingBoxes MaxPool1d MaxPool2d MaxPool3d MaxPooling1D MaxPooling2D MaxPooling3D MaxUnpool1d MaxUnpool2d MaxUnpool3d Maximum Maxpool Mean Median MelSpectrogram Meshgrid Mgrid Min Minimum Mish MixUp Mod Modf Module ModuleList Moveaxis Mse Mul MultiSteps MultiStepsState MultiTransformState MultiheadAttention Multinomial Multiply NONE Names Nan NanArgMax NanArgMin NanCumprod NanCumsum NanMax NanMean NanMedian NanMin NanPercentile NanProd NanToNum Nll NonNegativeParamsState NonZero Nop Normalization Normalize Object OneHot Ones OptState OptaxTest Optimizer Optional PRelu Pad Param Parameter ParameterList PartitionState Permute Pipeline Poly Polyint Polymul Polysub Polyval Positive Pow Power Printoptions Prod PromoteTypes Ptp Put PutAlongAxis Quantile RMSNormalization RMSprop RNN RNNBase RPC_AVAILABLE R_ Rad2deg Radians RandAugment RandInt RandomBrightness RandomColorDegeneration RandomColorJitter RandomContrast RandomCrop RandomElasticTransform RandomErasing RandomFlip RandomGaussianBlur RandomGrayscale RandomHorizontalFlip RandomHue RandomInvert RandomPerspective RandomPosterization RandomRotation RandomSaturation RandomSharpness RandomShear RandomTranslation RandomVerticalFlip RandomZoom Ravel RavelMultiIndex ReLU Real Reciprocal Relu6 Remainder Repeat RepeatVector Rescaling Reshape Resize Resizing ResultType ReversibleEmbedding RightShift Rint Rmsnorm Rnncellbase Roll Rollaxis Roots Rot90 Round SGD STFT STFTSpectrogram STRING Save Savez ScaleByTrustRatioState ScaleState Scan Scatter SearchSorted Select Selu SeparableConv1D SeparableConv2D SeparableConvolution1D SeparableConvolution2D Sequential SetDiff1d SetPrintoptions SetXor1d Shape ShouldSkipUpdateFunction SiLU Sigmoid Sign Signbit SignedInteger SimpleRNN SimpleRNNCell Sin Sinc Single Sinh Size Slice Smoothl1 SnapshotState Softmax Softmax2d Softmin Softplus Softshrink Softsign Solarization Sort SortComplex Spacing Sparsemax SpatialDropout1D SpatialDropout2D SpatialDropout3D SpectralNorm SpectralNormalization Split Sqrt Square Squeeze Stack StackedRNNCells Std StepLR StringLookup Sub Subtract Sum Swapaxes Synchronize T TFSMLayer TUPLE T_module Take TakeAlongAxis Tan Tanh Tanhshrink Tensor TensorType Tensordot TextVectorization Threshold Tile TimeDistributed ToTensor TopK TorchFunctional TorchLoad TorchModuleWrapper TorchSave Trace TraceState Transformer TransformerDecoderLayer TransformerEncoderLayer Transformerdecoder Transformerencoder Transpose Trapezoid Tri Tril TrilIndices TrilIndicesFrom TrimZeros Triu TriuIndices TriuIndicesFrom TrueDivide Trunc UInt8 Ufunc Uint Uint16 Uint2 Uint32 Uint4 Uint64 Unflatten Unfold Union1d Unique UniqueAll UniqueCounts UniqueInverse UniqueValues UnitNormalization UnpackBits UnravelIndex UnsignedInteger Unsqueeze Unstack Unwrap UpSampling1D UpSampling2D UpSampling3D UpdateCvState Upsample V Vander Var Variable Vdot Vecdot Vecmat Vectorize View Vsplit Vstack WeightNorm Where WrappedSchedule Wrapper ZeroNansState ZeroPadding1D ZeroPadding2D ZeroPadding3D Zeros ZerosLike __array_namespace_info__ absolute activation activation_relu_or_gelu adaptive_average_pool adaptive_grad_clip adaptive_max_pool add_decayed_weights add_noise add_scale addressable_data align_corners alpha amax amin apply apply_along_axis apply_gradients apply_if_finite apply_mask apply_over_axes approximate arccos arccosh arcsin arcsinh arctan arctan2 arctanh args array_equal array_equiv array_split as_string assert_equal associative_scan at atleast_1d atleast_2d atleast_3d attr average_pool axis base batch_normalization batch_shape beta bias bias_correction bias_hh bias_ih bias_k bias_v binary_cross_entropy binomial bits bitwise_and bitwise_invert bitwise_left_shift bitwise_not bitwise_or bitwise_right_shift bitwise_xor bool_ broadcast broadcast_arrays broadcast_shapes broadcast_to byte cached can_cast cast_like categorical categorical_crossentropy cdouble ceil_mode chain char character children chisquare choice cholesky cholesky_inverse clear clip_by_block_rms clip_by_global_norm clip_grad_norm close column_stack complex_ complexfloating compute_mask cond conditionally_mask conditionally_transform constant_ constant_schedule control_delta_method control_variates_jacobians conv1d_input conv1d_weight conv2d_input conv2d_weight conv3d_input conv3d_weight conv_transpose conv_transpose1d conv_transpose2d conv_transpose3d convert_to_numpy convert_to_tensor convex_kl_divergence copy_to_host_async cosine_decay_schedule cosine_distance cosine_onecycle_schedule cosine_similarity count count_nonzero cpu create_mask cross_entropy ctc_decode ctc_loss ctc_loss_with_forward_probs cumulative_prod cumulative_sum custom_gradient d_model data data_ptr decoder depthwise_conv deserialize device_ids device_mesh device_type diag_indices diag_indices_from digamma dilation dim dims dirichlet div_value divide_no_nan dot_product_attention downscale_factor dropout1 dropout2 dropout3 eig eigh eigvals eigvalsh einsum_path ema embedding_dim empty_like enable_grad encoder end_dim eps erfc erfcx erfinv estimate_control_variate_coefficients euler_gamma eval expand_dims expit exponential exponential_decay extend extract_sequences eye_ f fft fft2 fftfreq fftn fftshift fill_diagonal filters find_duplicates flatten_parameters float floor_divide fn fori_loop forward freeze from_dlpack full_like gamma gammainc gammaincc gammaln gaussian_nll_loss geometric get get_all_with_path get_item get_parameter get_printoptions glorot_normal glorot_uniform grad greater_equal group_size groups gumbel half hard_shrink hard_sigmoid hard_silu hard_swish hard_tanh he_normal he_uniform head head_dim hfft hidden_size hinge_loss histogram_bin_edges huber_loss ifft ifft2 ifftn ifftshift ihfft iinfo in1_features in2_features in_features in_top_k index inexact inf init inject_hyperparams inject_stateful_hyperparams inplace int int2 int4 int_ interpolate inv invert_permutation irfft irfft2 irfftn is_tensor istft item items itemsize jit join join_schedules jvp k kaiming_normal kaiming_uniform keep_params_nonnegative kernel_size key keys kl_divergence l1_loss l2_loss label_smoothing lambd laplace layer_norm layers leaky_relu left_shift less_equal linear1 linear2 linear_onecycle_schedule linear_schedule linear_to_mel_weight_matrix load_state_dict log_cosh log_ndtr log_prob log_sigmoid log_softmax logdet logger logical_and logical_not logical_or logical_xor logistic logit lognormal long lower lstsq lu lu_factor map map_params margin margin_ranking_loss mask mask_check mask_indices mask_mod masked matrix_norm matrix_power matrix_rank matrix_transpose max_norm max_pool max_pool1d max_pool2d max_pool3d max_val measure_valued_estimation_mean measure_valued_estimation_std measure_valued_jacobians merge_masks metadata min_val mode modules moments momentum moving_avg_baseline mse_loss multi_dot multi_transform multivariate_normal n name named_chain named_modules named_parameters nan_to_num nbytes ndarray ndim negative negative_slope nextafter nll_loss no_grad nonlinearity norm norm1 norm2 norm3 norm_first norm_type normal not_equal ntxent num_embeddings num_features num_groups num_heads num_layers number object_ one_hot ones_ ones_like orthogonal orthogonal_ out_features out_proj outer output_device output_ratio output_size p padding padding_idx parameters parameters_to_vector params pareto partition pathwise_jacobians per_example_global_norm_clip per_example_layer_norm_clip permutation permute_dims pi piecewise piecewise_constant_schedule piecewise_interpolate_schedule pinv place poisson polar polydiv polynomial_schedule pop popitem predict promote_types psi psnr qr quantize randn random_like rank ravel_multi_index rayleigh register_buffer register_parameter remove remove_axis replace reset reshape_weight_to_matrix result_type rfft rfft2 rfftfreq rfftn right_shift rms_norm rms_normalization rsqrt safe_softmax_cross_entropy saturate_cast savez_compressed scale scale_factor scale_grad_by_freq scaled_dot_product_attention scatter_update score_function_jacobians seed segment_max segment_sum selective_transform separable_conv serialize set set_printoptions sgdr_schedule short shuffle sigmoid_binary_cross_entropy sigmoid_focal_loss skip_large_updates skip_not_finite slice_update slogdet smooth_l1_loss smooth_labels snapshot soft_shrink softmax_cross_entropy softmax_cross_entropy_with_integer_labels solve solve_triangular sparse sparse_ sparse_categorical_crossentropy sparse_plus sparse_sigmoid split_key_like squared_error squareplus start_dim state state_dict static_graph step stop_gradient stride strides svd svdvals swap switch take_along_axis tanh_shrink tensorinv tensorsolve test_import threshold_ to_dense to_device tolist top_k train training tree_add_scalar_mul tree_l1_norm tree_l2_norm tree_linf_norm tree_map_params tree_scalar_mul triangular tril_indices tril_indices_from trim_zeros triu_indices triu_indices_from true_divide trunc_normal_ truncated_normal type unflattened_size uniform unique_all unique_counts unique_inverse unique_values unravel_index unwrap_random_key_data update update_infinity_moment update_moment update_moment_per_elem_norm upper upsample_bilinear upsample_nearest upscale_factor value value_and_grad values variant vector_norm vectorized_map view_as_complex view_as_real vmap wald warmup_constant_schedule warmup_cosine_decay_schedule warmup_exponential_decay_schedule weight weight_hh weight_ih while_loop xavier_normal xavier_uniform xlog1py zero_grad zero_nans zeros_ zeros_like