fit_a_nef.nef.utils.custom_uniform

fit_a_nef.nef.utils.custom_uniform(numerator: float = 6, mode: str = 'fan_in', dtype: ~numpy.dtype = <class 'jax.numpy.float64'>, in_axis: int | ~typing.Sequence[int] = -2, out_axis: int | ~typing.Sequence[int] = -1, batch_axis: ~typing.Sequence[int] = (), distribution: str = 'uniform') → Initializer

Builds an initializer that returns real uniformly-distributed random arrays.

Parameters:

• numerator (float) – the numerator of the range of the random distribution.

• mode (str) – the mode for computing the range of the random distribution.

• dtype (jnp.dtype) – optional; the initializer’s default dtype.

• in_axis (Union[int, Sequence[int]]) – the axis or axes that specify the input size.

• out_axis (Union[int, Sequence[int]]) – the axis or axes that specify the output size.

• batch_axis (Sequence[int]) – the axis or axes that specify the batch size.

• distribution (str) – the distribution of the random distribution.

Returns:

An initializer that returns arrays whose values are uniformly distributed in the range [-range, range).

Return type:

Initializer