deepfold.utils package

Subpackages

Submodules

deepfold.utils.checkpoint_utils module

deepfold.utils.checkpoint_utils.map_init_state_dicts(alphafold_state_dict_keys: List[str], init_alphafold_state_dict: dict, init_optimizer_state_dict: dict) → Tuple[dict, dict][source]
deepfold.utils.checkpoint_utils.resume_checkpoint(alphafold: AlphaFold, optimizer: Optimizer | None, swa_alphafold: AlphaFoldSWA | None, checkpoint_dirpath: Path, device: device) → int[source]
deepfold.utils.checkpoint_utils.resume_from_latest_checkpoint(alphafold: AlphaFold, optimizer: Optimizer | None, swa_alphafold: AlphaFoldSWA | None, checkpoints_dirpath: Path, device: device) → int[source]
deepfold.utils.checkpoint_utils.save_checkpoint(alphafold: AlphaFold, optimizer: Optimizer, swa_alphafold: AlphaFoldSWA, iteration: int, checkpoint_dirpath: Path) → None[source]
deepfold.utils.checkpoint_utils.save_last_checkpoint(alphafold: AlphaFold, optimizer: Optimizer, swa_alphafold: AlphaFoldSWA, iteration: int, checkpoints_dirpath: Path, keep_last_checkpoints: int) → None[source]
deepfold.utils.checkpoint_utils.save_val_checkpoint(alphafold: AlphaFold | AlphaFoldSWA, iteration: int, checkpoints_dirpath: Path, keep_val_checkpoints: bool, val_avg_lddt_ca: float) → None[source]

deepfold.utils.chunk_utils module

deepfold.utils.chunk_utils.chunk_layer(layer: Callable, inputs: Dict[str, Any], chunk_size: int, num_batch_dims: int, low_mem: bool = False, _out: Any | None = None, _add_into_out: bool = False) → Any[source]

Implements the “chunking” procedure described in section 1.11.8.

Layer outputs and inputs are assumed to be simple “pytrees,” consisting only of (arbitrarily nested) lists, tuples, and dicts with torch.Tensor leaves.

Parameters:

  • layer – The layer to be applied chunk-wise

  • inputs – A (non-nested) dictionary of keyworded inputs. All leaves must be tensors and must share the same batch dimensions.

  • chunk_size – The number of sub-batches per chunk. If multiple batch dimensions are specified, a “sub-batch” is defined as a single indexing of all batch dimensions simultaneously (s.t. the number of sub-batches is the product of the batch dimensions).

  • no_batch_dims – How many of the initial dimensions of each input tensor can be considered batch dimensions.

  • low_mem – Avoids flattening potentially large input tensors. Unnecessary in most cases, and is ever so slightly slower than the default setting.

Returns:

The reassembled output of the layer on the inputs.

deepfold.utils.crop_utils module

deepfold.utils.crop_utils.unpad_to_schema_shape_(outputs: Dict[str, ndarray], output_schema_shapes: Dict[str, tuple], num_residues: int, num_clustered_msa_seq: int) → Dict[str, ndarray][source]

deepfold.utils.datetime_utils module

deepfold.utils.datetime_utils.datetime_from_string(datetime_string: str, datetime_format: str = '%Y-%m-%d %H:%M:%S') → datetime[source]

Converts string to datetime object.

deepfold.utils.datetime_utils.datetime_to_string(datetime_object: datetime, string_format: str = '%Y-%m-%d %H:%M:%S') → str[source]

Converts datetime object to string.

deepfold.utils.datetime_utils.get_timestamp_string() → str[source]

Returns timestamp in YYYYmmdd_HHMMSS_ffffff format.

deepfold.utils.dist_utils module

deepfold.utils.dist_utils.get_pad_size(tensor: Tensor, dim: int, num_chunks: int) → int[source]
deepfold.utils.dist_utils.pad_tensor(tensor: Tensor, dim: int, pad_size: int) → Tensor[source]

deepfold.utils.feats_utils module

deepfold.utils.feats_utils.crop_a3m_string(a3m_string: str, start: int, end: int) → str[source]
deepfold.utils.feats_utils.crop_features(feats: dict, start: int, end: int)[source]
Parameters:

  • [start (start - 1, end)

  • range (end] ->)

deepfold.utils.file_utils module

deepfold.utils.file_utils.dump_pickle(obj: Any, path: PathLike, level: int = 5) → None[source]
deepfold.utils.file_utils.find_paths(paths: Sequence[PathLike]) → List[Path][source]
deepfold.utils.file_utils.get_file_content_and_extension(file_path: PathLike) → Tuple[str, str][source]

Reads the content of a text file or compressed text file and returns its content along with the real extension (excluding compression extensions like .gz).

Parameters:

file_path (str) – Path to the text file or compressed text file.

Returns:

A tuple containing the file content (str) and the real extension (str).

Return type:

tuple

deepfold.utils.file_utils.load_pickle(path: PathLike) → Any[source]
deepfold.utils.file_utils.read_text(path: PathLike) → str[source]
deepfold.utils.file_utils.restore_wrapper(func)[source]

deepfold.utils.hash_utils module

deepfold.utils.hash_utils.hash_string_into_number(s: str) → int[source]

Hashes string into uint64-like integer number.

deepfold.utils.import_utils module

class deepfold.utils.import_utils.Param(param: Union[torch.Tensor, List[torch.Tensor]], param_type: deepfold.utils.import_utils.ParamType = <ParamType.Other: functools.partial(<function ParamType.<lambda> at 0x7fc7f62e5090>)>, stacked: bool = False, swap: bool = False)[source]

Bases: object

param: Tensor | List[Tensor]
param_type: ParamType = functools.partial(<function ParamType.<lambda>>)
stacked: bool = False
swap: bool = False
class deepfold.utils.import_utils.ParamType(value)[source]

Bases: Enum

An enumeration.

LinearBiasMHA = functools.partial(<function ParamType.<lambda>>)
LinearMHAOutputWeight = functools.partial(<function ParamType.<lambda>>)
LinearWeight = functools.partial(<function ParamType.<lambda>>)
LinearWeightMHA = functools.partial(<function ParamType.<lambda>>)
LinearWeightOPM = functools.partial(<function ParamType.<lambda>>)
Other = functools.partial(<function ParamType.<lambda>>)
deepfold.utils.import_utils.assign(translation_dict, orig_weights)[source]
deepfold.utils.import_utils.import_jax_weights_(model: AlphaFold, npz_path: str, is_multimer: bool = False, enable_ptm: bool = False, enable_templates: bool = False, fuse_projection_weights: bool = False) → None[source]

Import AlphaFold JAX parameters.

Parameters:

  • model – AlphaFold model.

  • npz_path – Path to an NPZ file.

  • is_multimer – Whether multimer model or not.

  • enable_ptm – Enable predicted aligned error related modules.

  • enable_templates – Enable template related modules.

  • fuse_projection_weights – Whether triangular multiplicative layers are fused or not.

deepfold.utils.import_utils.reshape_weight(x: ndarray) → ndarray[source]
deepfold.utils.import_utils.stacked(param_dict_list, out=None)[source]
Parameters:

param_dict_list – A list of (nested) Param dicts to stack. The structure of each dict must be the identical (down to the ParamTypes of “parallel” Params). There must be at least one dict in the list.

deepfold.utils.iter_utils module

deepfold.utils.iter_utils.all_equal(values: list) → bool[source]

Checks if all values in list are equal to each other.

deepfold.utils.iter_utils.flatten_dict(d: MutableMapping, parent_key: str = '', sep: str = '.')[source]
deepfold.utils.iter_utils.list_zip(*arglists) → list[source]

Transforms given columns into list of rows.

deepfold.utils.iter_utils.map_dict_values(fn: Callable, d: dict) → dict[source]

Maps dictionary values using given function.

deepfold.utils.iter_utils.slice_generator(start: int, end: int, size: int) → Iterator[Tuple[int, int]][source]

Returns slice indices iterator from start to end.

deepfold.utils.log_utils module

class deepfold.utils.log_utils.TqdmHandler[source]

Bases: StreamHandler

emit(record: LogRecord) → None[source]

Emit a record.

If a formatter is specified, it is used to format the record. The record is then written to the stream with a trailing newline. If exception information is present, it is formatted using traceback.print_exception and appended to the stream. If the stream has an ‘encoding’ attribute, it is used to determine how to do the output to the stream.

deepfold.utils.log_utils.read_logs(filepath: Path, drop_overridden_iterations: bool = True) → DataFrame[source]
deepfold.utils.log_utils.save_logs(logs: List[dict], outpath: Path, append: bool) → None[source]
deepfold.utils.log_utils.setup_logging(filename: str | PathLike | Path, mode: str = 'w') → None[source]

deepfold.utils.precision module

deepfold.utils.precision.is_fp16_enabled() → bool[source]

deepfold.utils.random module

deepfold.utils.random.get_seed_from_string(s: str) → int[source]

Hashes input string and returns uint64-like integer seed value.

deepfold.utils.random.get_seed_randomly() → int[source]

Returns truly pseduorandom uint64-like integer seed value.

deepfold.utils.random.numpy_seed(seed: int | None = None, *additional_seeds, key: str = '')[source]

Context manager which seeds the Numpy PRNG with the seed and restores the state.

deepfold.utils.random.str_hash(string: str)[source]

deepfold.utils.rigid_utils module

class deepfold.utils.rigid_utils.Rigid(rots: Rotation | None, trans: Tensor | None)[source]

Bases: object

A class representing a rigid transformation. Little more than a wrapper around two objects: a Rotation object and a [*, 3] translation Designed to behave approximately like a single torch tensor with the shape of the shared batch dimensions of its component parts.

apply(pts: Tensor) → Tensor[source]

Applies the transformation to a coordinate tensor.

Parameters:

pts – A [*, 3] coordinate tensor.

Returns:

The transformed points.

apply_rot_fn(fn: Callable[[Rotation], Rotation]) → Rigid[source]

Applies a Rotation -> Rotation function to the stored rotation object.

Parameters:

fn – A function of type Rotation -> Rotation

Returns:

A transformation object with a transformed rotation.

apply_trans_fn(fn: Callable[[Tensor], Tensor]) → Rigid[source]

Applies a Tensor -> Tensor function to the stored translation.

Parameters:

fn – A function of type Tensor -> Tensor to be applied to the translation

Returns:

A transformation object with a transformed translation.

static cat(ts: Sequence[Rigid], dim: int) → Rigid[source]

Concatenates transformations along a new dimension.

Parameters:

  • ts – A list of T objects

  • dim – The dimension along which the transformations should be concatenated

Returns:

A concatenated transformation object

compose(r: Rigid) → Rigid[source]

Composes the current rigid object with another.

Parameters:

r – Another Rigid object

Returns:

The composition of the two transformations

compose_q_update_vec(q_update_vec: Tensor) → Rigid[source]

Composes the transformation with a quaternion update vector of shape [*, 6], where the final 6 columns represent the x, y, and z values of a quaternion of form (1, x, y, z) followed by a 3D translation.

Parameters:

q_vec – The quaternion update vector.

Returns:

The composed transformation.

cuda() → Rigid[source]

Moves the transformation object to GPU memory

Returns:

A version of the transformation on GPU

property device: device

Returns the device on which the Rigid’s tensors are located.

Returns:

The device on which the Rigid’s tensors are located

property dtype: dtype

Returns the dtype of the Rigid tensors.

Returns:

The dtype of the Rigid tensors

static from_3_points(p_neg_x_axis: Tensor, origin: Tensor, p_xy_plane: Tensor, eps: float = 1e-08) → Rigid[source]

Implements algorithm 21. Constructs transformations from sets of 3 points using the Gram-Schmidt algorithm.

Parameters:

  • p_neg_x_axis – [*, 3] coordinates

  • origin – [*, 3] coordinates used as frame origins

  • p_xy_plane – [*, 3] coordinates

  • eps – Small epsilon value

Returns:

A transformation object of shape [*]

static from_tensor_4x4(t: Tensor) → Rigid[source]

Constructs a transformation from a homogenous transformation tensor.

Parameters:

t – [*, 4, 4] homogenous transformation tensor

Returns:

T object with shape [*]

static from_tensor_7(t: Tensor, normalize_quats: bool = False) → Rigid[source]
get_rots() → Rotation[source]

Getter for the rotation.

Returns:

The rotation object

get_trans() → Tensor[source]

Getter for the translation.

Returns:

The stored translation

static identity(shape: Tuple[int], dtype: dtype | None = None, device: device | None = None, requires_grad: bool = True, fmt: str = 'quat') → Rigid[source]

Constructs an identity transformation.

Parameters:

  • shape – The desired shape

  • dtype – The dtype of both internal tensors

  • device – The device of both internal tensors

  • requires_grad – Whether grad should be enabled for the internal tensors

Returns:

The identity transformation

invert() → Rigid[source]

Inverts the transformation.

Returns:

The inverse transformation.

invert_apply(pts: Tensor) → Tensor[source]

Applies the inverse of the transformation to a coordinate tensor.

Parameters:

pts – A [*, 3] coordinate tensor

Returns:

The transformed points.

static make_transform_from_reference(n_xyz, ca_xyz, c_xyz, eps=1e-20)[source]

Returns a transformation object from reference coordinates.

Note that this method does not take care of symmetries. If you provide the atom positions in the non-standard way, the N atom will end up not at [-0.527250, 1.359329, 0.0] but instead at [-0.527250, -1.359329, 0.0]. You need to take care of such cases in your code.

Parameters:

  • n_xyz – A [*, 3] tensor of nitrogen xyz coordinates.

  • ca_xyz – A [*, 3] tensor of carbon alpha xyz coordinates.

  • c_xyz – A [*, 3] tensor of carbon xyz coordinates.

Returns:

A transformation object. After applying the translation and rotation to the reference backbone, the coordinates will approximately equal to the input coordinates.

map_tensor_fn(fn: Callable[[Tensor], Tensor]) → Rigid[source]

Apply a Tensor -> Tensor function to underlying translation and rotation tensors, mapping over the translation/rotation dimensions respectively.

Parameters:

fn – A Tensor -> Tensor function to be mapped over the Rigid

Returns:

The transformed Rigid object

scale_translation(trans_scale_factor: float) → Rigid[source]

Scales the translation by a constant factor.

Parameters:

trans_scale_factor – The constant factor

Returns:

A transformation object with a scaled translation.

property shape: Size

Returns the shape of the shared dimensions of the rotation and the translation.

Returns:

The shape of the transformation

stop_rot_gradient() → Rigid[source]

Detaches the underlying rotation object

Returns:

A transformation object with detached rotations

to_tensor_4x4() → Tensor[source]

Converts a transformation to a homogenous transformation tensor.

Returns:

A [*, 4, 4] homogenous transformation tensor

to_tensor_7() → Tensor[source]

Converts a transformation to a tensor with 7 final columns, four for the quaternion followed by three for the translation.

Returns:

A [*, 7] tensor representation of the transformation

unsqueeze(dim: int) → Rigid[source]

Analogous to torch.unsqueeze. The dimension is relative to the shared dimensions of the rotation/translation.

Parameters:

dim – A positive or negative dimension index.

Returns:

The unsqueezed transformation.

class deepfold.utils.rigid_utils.Rotation(rot_mats: Tensor | None = None, quats: Tensor | None = None, normalize_quats: bool = True)[source]

Bases: object

A 3D rotation. Depending on how the object is initialized, the rotation is represented by either a rotation matrix or a quaternion, though both formats are made available by helper functions. To simplify gradient computation, the underlying format of the rotation cannot be changed in-place. Like Rigid, the class is designed to mimic the behavior of a torch Tensor, almost as if each Rotation object were a tensor of rotations, in one format or another.

apply(pts: Tensor) → Tensor[source]

Apply the current Rotation as a rotation matrix to a set of 3D coordinates.

Parameters:

pts – A [*, 3] set of points

Returns:

[*, 3] rotated points

static cat(rs: Sequence[Rotation], dim: int) → Rotation[source]

Concatenates rotations along one of the batch dimensions. Analogous to torch.cat().

Note that the output of this operation is always a rotation matrix, regardless of the format of input rotations.

Parameters:

  • rs – A list of rotation objects

  • dim – The dimension along which the rotations should be concatenated

Returns:

A concatenated Rotation object in rotation matrix format

compose_q(r: Rotation, normalize_quats: bool = True) → Rotation[source]

Compose the quaternions of the current Rotation object with those of another.

Depending on whether either Rotation was initialized with quaternions, this function may call torch.linalg.eigh.

Parameters:

r – An update rotation object

Returns:

An updated rotation object

compose_q_update_vec(q_update_vec: Tensor, normalize_quats: bool = True) → Rotation[source]

Returns a new quaternion Rotation after updating the current object’s underlying rotation with a quaternion update, formatted as a [*, 3] tensor whose final three columns represent x, y, z such that (1, x, y, z) is the desired (not necessarily unit) quaternion update.

Parameters:

  • q_update_vec – A [*, 3] quaternion update tensor

  • normalize_quats – Whether to normalize the output quaternion

Returns:

An updated Rotation

compose_r(r: Rotation) → Rotation[source]

Compose the rotation matrices of the current Rotation object with those of another.

Parameters:

r – An update rotation object

Returns:

An updated rotation object

cuda() → Rotation[source]

Analogous to the cuda() method of torch Tensors

Returns:

A copy of the Rotation in CUDA memory

detach() → Rotation[source]

Returns a copy of the Rotation whose underlying Tensor has been detached from its torch graph.

Returns:

A copy of the Rotation whose underlying Tensor has been detached from its torch graph

property device: device

The device of the underlying rotation

Returns:

The device of the underlying rotation

property dtype: dtype

Returns the dtype of the underlying rotation.

Returns:

The dtype of the underlying rotation

get_cur_rot() → Tensor[source]

Return the underlying rotation in its current form

Returns:

The stored rotation

get_quats() → Tensor[source]

Returns the underlying rotation as a quaternion tensor.

Depending on whether the Rotation was initialized with a quaternion, this function may call torch.linalg.eigh.

Returns:

The rotation as a quaternion tensor.

get_rot_mats() → Tensor[source]

Returns the underlying rotation as a rotation matrix tensor.

Returns:

The rotation as a rotation matrix tensor

static identity(shape, dtype: dtype | None = None, device: device | None = None, requires_grad: bool = True, fmt: str = 'quat') → Rotation[source]

Returns an identity Rotation.

Parameters:

  • shape – The “shape” of the resulting Rotation object. See documentation for the shape property

  • dtype – The torch dtype for the rotation

  • device – The torch device for the new rotation

  • requires_grad – Whether the underlying tensors in the new rotation object should require gradient computation

  • fmt – One of “quat” or “rot_mat”. Determines the underlying format of the new object’s rotation

Returns:

A new identity rotation

invert() → Rotation[source]

Returns the inverse of the current Rotation.

Returns:

The inverse of the current Rotation

invert_apply(pts: Tensor) → Tensor[source]

The inverse of the apply() method.

Parameters:

pts – A [*, 3] set of points

Returns:

[*, 3] inverse-rotated points

map_tensor_fn(fn: Callable[[Tensor], Tensor]) → Rotation[source]

Apply a Tensor -> Tensor function to underlying rotation tensors, mapping over the rotation dimension(s). Can be used e.g. to sum out a one-hot batch dimension.

Parameters:

fn – A Tensor -> Tensor function to be mapped over the Rotation

Returns:

The transformed Rotation object

property requires_grad: bool

Returns the requires_grad property of the underlying rotation

Returns:

The requires_grad property of the underlying tensor

property shape: Size

Returns the virtual shape of the rotation object. This shape is defined as the batch dimensions of the underlying rotation matrix or quaternion. If the Rotation was initialized with a [10, 3, 3] rotation matrix tensor, for example, the resulting shape would be [10].

Returns:

The virtual shape of the rotation object

to(device: device | None, dtype: dtype | None) → Rotation[source]

Analogous to the to() method of torch Tensors

Parameters:

  • device – A torch device

  • dtype – A torch dtype

Returns:

A copy of the Rotation using the new device and dtype

unsqueeze(dim: int) → Rotation[source]

Analogous to torch.unsqueeze. The dimension is relative to the shape of the Rotation object.

Parameters:

dim – A positive or negative dimension index.

Returns:

The unsqueezed Rotation.

deepfold.utils.rigid_utils.identity_quats(batch_dims: Tuple[int], dtype: dtype | None = None, device: device | None = None, requires_grad: bool = True) → Tensor[source]
deepfold.utils.rigid_utils.identity_rot_mats(batch_dims: Tuple[int], dtype: dtype | None = None, device: device | None = None, requires_grad: bool = True) → Tensor[source]
deepfold.utils.rigid_utils.identity_trans(batch_dims: Tuple[int], dtype: dtype | None = None, device: device | None = None, requires_grad: bool = True) → Tensor[source]
deepfold.utils.rigid_utils.invert_quat(quat: Tensor)[source]
deepfold.utils.rigid_utils.invert_rot_mat(rot_mat: Tensor)[source]
deepfold.utils.rigid_utils.quat_multiply(quat1, quat2)[source]

Multiply a quaternion by another quaternion.

deepfold.utils.rigid_utils.quat_multiply_by_vec(quat, vec)[source]

Multiply a quaternion by a pure-vector quaternion.

deepfold.utils.rigid_utils.quat_to_rot(quat: Tensor) → Tensor[source]

Converts a quaternion to a rotation matrix.

Parameters:

quat – [*, 4] quaternions

Returns:

[*, 3, 3] rotation matrices

deepfold.utils.rigid_utils.rot_matmul(a: Tensor, b: Tensor) → Tensor[source]

Performs matrix multiplication of two rotation matrix tensors. Written out by hand to avoid AMP downcasting.

Parameters:

  • a – [*, 3, 3] left multiplicand

  • b – [*, 3, 3] right multiplicand

Returns:

The product ab

deepfold.utils.rigid_utils.rot_to_quat(rot: Tensor) → Tensor[source]
deepfold.utils.rigid_utils.rot_vec_mul(r: Tensor, t: Tensor) → Tensor[source]

Applies a rotation to a vector. Written out by hand to avoid transfer to avoid AMP downcasting.

Parameters:

  • r – [*, 3, 3] rotation matrices

  • t – [*, 3] coordinate tensors

Returns:

[*, 3] rotated coordinates

deepfold.utils.script_utils module

deepfold.utils.script_utils.get_basename(path: PathLike)[source]
deepfold.utils.script_utils.make_output_directory(output_dir: PathLike, model_name: str)[source]
deepfold.utils.script_utils.parse_fasta(fasta_str: str)[source]

deepfold.utils.swa_utils module

class deepfold.utils.swa_utils.AlphaFoldSWA(alphafold: AlphaFold, enabled: bool, decay_rate: float)[source]

Bases: Module

Weight averaging wrapper.

forward(batch)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

update(model: AlphaFold) → None[source]
class deepfold.utils.swa_utils.swa_avg_fn(decay_rate: float)[source]

Bases: object

Averaging function for EMA with configurable decay rate.

Suppl. ‘1.11.7 Evaluator setup’.

deepfold.utils.tensor_utils module

Utilities related to tensor operations.

deepfold.utils.tensor_utils.add(m1: Tensor, m2: Tensor, inplace: bool) → Tensor[source]
deepfold.utils.tensor_utils.array_tree_map(fn, tree, *, leaf_type=<class 'numpy.ndarray'>)
deepfold.utils.tensor_utils.batched_gather(data: Tensor, inds: Tensor, dim: int = 0, num_batch_dims: int = 0) → Tensor[source]
deepfold.utils.tensor_utils.collate(samples: List[dict]) → dict[source]

Converts list of samples into a batch dict.

deepfold.utils.tensor_utils.dict_map(fn: Callable[[T], Any], dic: Dict[Any, dict | list | tuple | T], leaf_type: Type[T]) → Dict[Any, dict | list | tuple | Any][source]
deepfold.utils.tensor_utils.dict_multimap(fn: Callable, dicts: List[Dict[Any, Any]]) → Dict[Any, Any][source]
deepfold.utils.tensor_utils.flatten_final_dims(t: Tensor, num_dims: int) → Tensor[source]
deepfold.utils.tensor_utils.masked_mean(mask: Tensor, value: Tensor, dim: int | Tuple[int, ...], eps: float = 0.0001, keepdim: bool = False) → Tensor[source]
deepfold.utils.tensor_utils.one_hot(x: Tensor, v_bins: Tensor) → Tensor[source]
deepfold.utils.tensor_utils.permute_final_dims(tensor: Tensor, inds: List[int]) → Tensor[source]
deepfold.utils.tensor_utils.pts_to_distogram(pts: Tensor, min_bin: float = 2.2325, max_bin: float = 21.6875, num_bins: int = 64) → Tensor[source]
deepfold.utils.tensor_utils.tensor_tree_map(fn, tree, *, leaf_type=<class 'torch.Tensor'>)
deepfold.utils.tensor_utils.tree_map(fn: Callable[[T], Any], tree: T, leaf_type: Type[T]) → Any[source]
deepfold.utils.tensor_utils.tree_map(fn: Callable[[T], Any], tree: dict, leaf_type: Type[T]) → dict
deepfold.utils.tensor_utils.tree_map(fn: Callable[[T], Any], tree: list, leaf_type: Type[T]) → list
deepfold.utils.tensor_utils.tree_map(fn: Callable[[T], Any], tree: tuple, leaf_type: Type[T]) → tuple

Module contents