jnp.matmul defaults to bfloat16, but this helper function doesn't.

Rescale the maximum saturation in `rgb` to be 1.

Compute PSNR given an MSE (we assume the maximum pixel value is 1).

The positional encoding used by the original NeRF paper.

Helper function to output density.

Render all the pixels of an image (in test mode). Args: render_fn: function, jit-ed render function. rays: a `Rays` pytree, the rays to be

Resampling. Args: rng: random number generator (or None for deterministic resampling). origins: Tensor, [..., 3], ray origins. directio

Rescales camera poses according to maximum x/y/z value.

jnp.cos() on a TPU may NaN out for large values.

jnp.sin() on a TPU may NaN out for large values.

Stratified sampling along the rays. Args: rng: random generator. If `None`, use deterministic sampling. origins: [..., 3], ray origins.

Save an image (probably a depthmap) to disk as a float32 TIFF.

Save an image (probably RGB) in [0, 1] to disk as a uint8 PNG.

Split data into shards for multiple devices along the first dimension.

A cyclic and uniform colormap, see http://basecase.org/env/on-rainbows.

Piecewise-Constant PDF sampling from sorted bins. Args: rng: random number generator (or None for `linspace` sampling). bins: [..., num_bin

Transfer data to devices (GPU/TPU).

Returns the tokenized representation of given input string(s) Parameters ---------- texts : Union[str, List[str]] An input s

One optimization step. Args: model: The linen model. config: The configuration. rng: jnp.ndarray, random number generator. state: u

Collect the sharded tensor to the shape before sharding.

Visualize the "cell" each coordinate lives in, and highlight its edges.

Visualize the coordinate of each point within its "cell".

Visualize the distance of each coordinate from the origin.

Visualizes depth maps.

Visualize fake normals of `depth` (optionally scaled to be isotropic).

Visualize a bundle of rays.

A wrapper around other visualizations for easy integration.

Volumetric Rendering Function. Args: rgbs: jnp.ndarray(float32), color, [batch_size, num_samples, 3] weights: jnp.ndarray(float32), weights

Compute the weighted percentile(s) of a single vector.