sage.core.torch

Filename : torch.py Description : Short description of the file

Created on 2026-01-23 09:16:49

__author__ = Narenraju Nagarajan __copyright__ = Copyright 2026, ProjectName __license__ = MIT Licence __version__ = 0.0.1 __maintainer__ = Narenraju Nagarajan __affiliation__ = N/A __email__ = N/A __status__ = [‘inProgress’, ‘Archived’, ‘inUsage’, ‘Debugging’]

GitHub Repository: NULL

Documentation: NULL

Functions

nudge_backward_(foo, max_limit[, nudge_factor])

In-place nudge to keep foo <= max_limit with tiny safety margin.

nudge_forward_(foo, min_limit[, nudge_factor])

In-place nudge to keep foo >= min_limit with tiny safety margin.

torch_grad(func, args[, argnums, create_graph])

PyTorch equivalent of jax.grad.

torch_value_and_grad(fn, inputs, *[, argnums, ...])

Torch equivalent of jax.value_and_grad(fn, argnums).

Module Contents

nudge_backward_(foo, max_limit, nudge_factor=1e-06)[source]

In-place nudge to keep foo <= max_limit with tiny safety margin. Modifies foo directly.

Parameters:
Return type:

None

nudge_forward_(foo, min_limit, nudge_factor=1e-06)[source]

In-place nudge to keep foo >= min_limit with tiny safety margin. Modifies foo directly.

Parameters:
Return type:

None

torch_grad(func, args, argnums=0, create_graph=False)[source]

PyTorch equivalent of jax.grad.

Computes gradient of func w.r.t. args[argnums].

Parameters:

  • func – callable

  • args – tuple of arguments passed to func

  • argnums – int or tuple of ints (default: 0)

  • create_graph – whether to construct graph for higher-order grads

Returns:

Gradient(s) corresponding to argnums - Tensor if single argnum - Tuple of tensors if multiple argnums

torch_value_and_grad(fn, inputs, *, argnums=0, create_graph=False)[source]

Torch equivalent of jax.value_and_grad(fn, argnums).

Call style:

value, grads = torch_value_and_grad(fn, (arg1, arg2, arg3)) value, grads = torch_value_and_grad(fn, (arg1, arg2, arg3), argnums=(0, 1))

Default:

Differentiates w.r.t. argnums=0 (first argument only), matching JAX behavior.

Args:

fn: callable inputs: Tensor or tuple of Tensors argnums: int or tuple of ints specifying which arguments to differentiate create_graph: whether to construct higher-order graph

Returns:

:

value: fn(*inputs) grads: gradient(s) w.r.t. argnums

  • Tensor if argnums is int

  • Tuple[Tensor, …] if argnums is tuple

Parameters: