Skip to content

supaernova.steps.nflow.tf

source package supaernova.steps.nflow.tf

Classes

source class TFNFlowModel(config: NFlowModelStep[Literal['tf']], *args: Any, **kwargs: Any)

Bases : ks.Model

Attributes

  • name — Name of the layer (string), set in the constructor.

  • name_scope — Returns a tf.name_scope instance for this class.

  • variables — Returns the list of all layer variables/weights.

  • submodules — Sequence of all sub-modules.

  • dtype — The dtype of the layer weights.

  • supports_masking — Whether this layer supports computing a mask using compute_mask.

  • dynamic — Whether the layer is dynamic (eager-only); set in the constructor.

  • activity_regularizer — Optional regularizer function for the output of this layer.

  • input_spec — InputSpec instance(s) describing the input format for this layer.

  • weights — Returns the list of all layer variables/weights.

  • losses — List of losses added using the add_loss() API.

  • metrics — Return metrics added using compile() or add_metric().

  • input_mask — Retrieves the input mask tensor(s) of a layer.

  • output_mask — Retrieves the output mask tensor(s) of a layer.

  • input — Retrieves the input tensor(s) of a layer.

  • output — Retrieves the output tensor(s) of a layer.

  • input_shape — Retrieves the input shape(s) of a layer.

  • output_shape — Retrieves the output shape(s) of a layer.

  • dtype_policy — The dtype policy associated with this layer.

  • compute_dtype — The dtype of the layer's computations.

  • variable_dtype — Alias of Layer.dtype, the dtype of the weights.

  • inbound_nodes — Return Functional API nodes upstream of this layer.

  • outbound_nodes — Return Functional API nodes downstream of this layer.

  • metrics_names — Returns the model's display labels for all outputs.

  • distribute_strategy — The tf.distribute.Strategy this model was created under.

  • run_eagerly — Settable attribute indicating whether the model should run eagerly.

  • autotune_steps_per_execution — Settable property to enable tuning for steps_per_execution

  • steps_per_execution — Settable `steps_per_execution variable. Requires a compiled model.

  • jit_compile — Specify whether to compile the model with XLA.

  • distribute_reduction_method — The method employed to reduce per-replica values during training.

  • state_updates — Deprecated, do NOT use!

Methods

source method TFNFlowModel.build(input_shape) → None

source method TFNFlowModel.call(inputs: NFlowInputs, training: bool | None = None, mask: TensorCompatible | None = None) → NFlowOutputs

source method TFNFlowModel.u_to_z(inputs: NFlowInputs, *, permute: bool = False) → NFlowInputs

source method TFNFlowModel.z_to_u(inputs: NFlowInputs, *, permute: bool = False) → NFlowInputs

source method TFNFlowModel.z_to_u_steps(inputs: NFlowInputs, step: int, *, permute: bool = False) → tuple['NFlowInputs', bool]

source method TFNFlowModel.train_model(*, savepath: Path | None = None) → ks.callbacks.History

source method TFNFlowModel.get_data(latents: FTensor[S['batch_dim nspec_dim pae_latents']]) → FTensor[S['batch_dim nspec_dim n_flow_latents']]

source method TFNFlowModel.build_model() → None

source method TFNFlowModel.save_checkpoint(savepath: Path) → None

source method TFNFlowModel.load_checkpoint(loadpath: Path) → None

source method TFNFlowModel.get_config() → dict[str, 'Any']

source classmethod TFNFlowModel.from_config(config: dict[str, 'Any']) → Self

source method TFNFlowModel.build_from_config(_config: dict[str, 'Any']) → None

source method TFNFlowModel.set_seed(seed: int = 0) → None