autocast.types.batch

class Sample(input_fields, output_fields, constant_scalars, constant_fields, boundary_conditions)

Bases: object

A batch in input data space.

Parameters:

• input_fields (Float[Tensor, 'time spatial *spatial channel'])

• output_fields (Float[Tensor, 'time spatial *spatial channel'])

• constant_scalars (Float[Tensor, 'channel'] | None)

• constant_fields (Float[Tensor, 'spatial *spatial channel'] | None)

• boundary_conditions (Float[Tensor, 'spatial *spatial'] | None)

input_fields: Float[Tensor, 'time spatial *spatial channel']

output_fields: Float[Tensor, 'time spatial *spatial channel']

constant_scalars: Float[Tensor, 'channel'] | None

constant_fields: Float[Tensor, 'spatial *spatial channel'] | None

boundary_conditions: Float[Tensor, 'spatial *spatial'] | None

class EncodedSample(encoded_inputs, encoded_output_fields, global_cond, encoded_info)

Bases: object

A batch after being processed by an Encoder.

Parameters:

• encoded_inputs (Float[Tensor, 'batch *optional_dims channel'])

• encoded_output_fields (Float[Tensor, 'batch *optional_dims channel'])

• global_cond (Float[Tensor, '*optional_dims channel'] | None)

• encoded_info (dict[str, Tensor])

encoded_inputs: Float[Tensor, 'batch *optional_dims channel']

encoded_output_fields: Float[Tensor, 'batch *optional_dims channel']

global_cond: Float[Tensor, '*optional_dims channel'] | None

encoded_info: dict[str, Tensor]

class Batch(input_fields, output_fields, constant_scalars, constant_fields, boundary_conditions=None)

Bases: object

A batch in input data space.

Parameters:

• input_fields (Float[Tensor, 'batch time spatial *spatial channel'])

• output_fields (Float[Tensor, 'batch time spatial *spatial channel'])

• constant_scalars (Float[Tensor, 'batch channel'] | None)

• constant_fields (Float[Tensor, 'batch spatial *spatial channel'] | None)

• boundary_conditions (Float[Tensor, 'spatial *spatial'] | None)

input_fields: Float[Tensor, 'batch time spatial *spatial channel']

output_fields: Float[Tensor, 'batch time spatial *spatial channel']

constant_scalars: Float[Tensor, 'batch channel'] | None

constant_fields: Float[Tensor, 'batch spatial *spatial channel'] | None

boundary_conditions: Float[Tensor, 'spatial *spatial'] | None = None

repeat(m)

Repeat batch members.

This interleaves the batch dimension by repeating each sample m times.

For example, for m=3, a batch with samples 0, 1, 2, … becomes 0, 0, 0, 1, 1, 1, 2, 2, 2, …

Parameters:

m (int)

Return type:

Batch

to(device)

Move batch to device.

Parameters:

device (device | str)

Return type:

Batch

class EncodedBatch(encoded_inputs, encoded_output_fields, global_cond, encoded_info)

Bases: object

A batch after being processed by an Encoder.

Parameters:

• encoded_inputs (Float[Tensor, 'batch *optional_dims channel'])

• encoded_output_fields (Float[Tensor, 'batch *optional_dims channel'])

• global_cond (Float[Tensor, 'batch *optional_dims channel'] | None)

• encoded_info (dict[str, Tensor])

encoded_inputs: Float[Tensor, 'batch *optional_dims channel']

encoded_output_fields: Float[Tensor, 'batch *optional_dims channel']

global_cond: Float[Tensor, 'batch *optional_dims channel'] | None

encoded_info: dict[str, Tensor]

repeat(m)

Repeat batch members.

This interleaves the batch dimension by repeating each sample m times.

For example, for m=3, a batch with samples 0, 1, 2, … becomes 0, 0, 0, 1, 1, 1, 2, 2, 2, …

Parameters:

m (int)

Return type:

EncodedBatch

to(device)

Move batch to device.

Parameters:

device (device | str)

Return type:

EncodedBatch