Skip to content
 
ml-switcheroo ScaleState
Type to start searching
    ml-switcheroo
    • Operation Reference
    ml-switcheroo
    • Documentation
    • ml-switcheroo πŸ”„πŸ¦˜
    • Architecture
    • Extending
    • Extending with DSL
    • Maintenance
    • Internal Architecture & Theoretical Mechanics
    • Ideas
    • Reference
    • API Reference
    • Operation Reference
      • 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
        • Show Source
      • 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
      • 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
    • Show Source

    ScaleStateΒΆ

    Auto-generated from jax_code_defs

    JAX (Core)

    API: optax.ScaleState
    Strategy: Direct Mapping
    Official Docs β†—

    Flax NNX

    API: optax.ScaleState
    Strategy: Direct Mapping
    Official Docs β†—
    "Previous" ScaleByTrustRatioState
    "Next" Scan
    © Copyright 2026, Samuel Marks.
    Created using Sphinx 9.1.0. and Material for Sphinx