Skip to main content
`Ctrl`+`K`

# brevitas.ops package

## Submodules

## brevitas.ops.autograd_ste_ops module

Implementation of various torch.autograd.Function with straight-through estimators.

class brevitas.ops.autograd_ste_ops.AbsBinarySignGradFn(*args, **kwargs)

  Bases: `Function`

  Autograd function that implements `torch.abs()` with a binary-sign backward, in order to have subgradient 1 in 0. Compare with `torch.abs()`’ subgradient of 0 in 0.

  `AbsBinarySignGradFn.apply(*args)` is first aliased to `abs_binary_sign_grad(*args)` and then wrapped by `abs_binary_sign_grad()` when env `BREVITAS_JIT=0`. See `abs_binary_sign_grad()` for details on the interface and examples.

class brevitas.ops.autograd_ste_ops.BinarySignSteFn(*args, **kwargs)

  Bases: `Function`

  Autograd function that implements `binary_sign()` with a straight-through gradient estimator.

  `BinarySignSteFn.apply(*args)` is first aliased to `binary_sign_ste_impl(*args)` and then wrapped by `binary_sign_ste()` when env `BREVITAS_JIT=0`. See `binary_sign_ste()` for details on the interface and examples.

class brevitas.ops.autograd_ste_ops.CeilSteFn(*args, **kwargs)

  Bases: `Function`

  Autograd function that implements `torch.ceil()` with a straight-through gradient estimator.

  `CeilSteFn.apply(*args)` is first aliased to `ceil_ste_impl(*args)` and then wrapped by `ceil_ste()` when env `BREVITAS_JIT=0`. See `ceil_ste()` for details on the interface and examples.

class brevitas.ops.autograd_ste_ops.DPURoundSteFn(*args, **kwargs)

  Bases: `Function`

  Autograd function that implements `dpu_round()` with a straight-through gradient estimator.

  `DPURoundSteFn.apply(*args)` is first aliased to `dpu_round_ste_impl(*args)` and then wrapped by `dpu_round_ste()` when env `BREVITAS_JIT=0`. See `dpu_round_ste()` for details on the interface and examples.

class brevitas.ops.autograd_ste_ops.FloorSteFn(*args, **kwargs)

  Bases: `Function`

  Autograd function that implements `torch.floor()` with a straight-through gradient estimator.

  `FloorSteFn.apply(*args)` is first aliased to `floor_ste_impl(*args)` and then wrapped by `floor_ste()` when env `BREVITAS_JIT=0`. See `floor_ste()` for details on the interface and examples.

class brevitas.ops.autograd_ste_ops.InplaceTensorClampSteFn(*args, **kwargs)

  Bases: `Function`

  Autograd function that implements `tensor_clamp_()` with a straight-through gradient estimator for the gradient of y w.r.t. to x, while the gradient of y w.r.t. to min_val and max_val is always None.

  `InplaceTensorClampSteFn.apply(*args)` is first aliased to `tensor_clamp_ste_impl_(*args)` and then wrapped by `tensor_clamp_()` when env `BREVITAS_JIT=0`. See `tensor_clamp_()` for details on the interface and examples.

class brevitas.ops.autograd_ste_ops.RoundSteFn(*args, **kwargs)

  Bases: `Function`

  Autograd function that implements `torch.round()` with a straight-through gradient estimator.

  `RoundSteFn.apply(*args)` is first aliased to `round_ste_impl(*args)` and then wrapped by `round_ste()` when env `BREVITAS_JIT=0`. See `round_ste()` for details on the interface and examples.

class brevitas.ops.autograd_ste_ops.RoundToZeroSteFn(*args, **kwargs)

  Bases: `Function`

  Autograd function that implements `round_to_zero()` with a straight-through gradient estimator.

  `RoundToZeroSteFn.apply(*args)` is first aliased to `round_to_zero_ste_impl(*args)` and then wrapped by `round_to_zero_ste()` when env `BREVITAS_JIT=0`. See `round_to_zero_ste()` for details on the interface and examples.

class brevitas.ops.autograd_ste_ops.ScalarClampMinSteFn(*args, **kwargs)

  Bases: `Function`

  Autograd function that implements `torch.clamp_min` with a straight-through gradient estimator for the gradient of y w.r.t. to x, while the gradient of y w.r.t. to `min_val` is always `None`.

  `ScalarClampMinSteFn.apply(*args)` is first aliased to `scalar_clamp_min_ste_impl(*args)` and then wrapped by `scalar_clamp_min_ste()` and invoked when env `BREVITAS_JIT=0`. See `scalar_clamp_ste()` for details on the interface and examples.

class brevitas.ops.autograd_ste_ops.ScalarClampSteFn(*args, **kwargs)

  Bases: `Function`

  Autograd function that implements `torch.clamp` with a straight-through gradient estimator for the gradient of y w.r.t. to x, while the gradient of y w.r.t. to `min_val` and `min_val` are always `None`.

  `ScalarClampSteFn.apply(*args)` is first aliased to `scalar_clamp_ste_impl(*args)` and then wrapped by `scalar_clamp_ste()` and invoked when env `BREVITAS_JIT=0`. See `scalar_clamp_ste()` for details on the interface and examples.

class brevitas.ops.autograd_ste_ops.TensorClampSteFn(*args, **kwargs)

  Bases: `Function`

  Autograd function that implements `tensor_clamp()` with a straight-through gradient estimator for the gradient of y w.r.t. to x, while the gradient of y w.r.t. to min_val and max_val is always None.

  `TensorClampSteFn.apply(*args)` is first aliased to `tensor_clamp_ste_impl(*args)` and then wrapped by `tensor_clamp()` when env `BREVITAS_JIT=0`. See `tensor_clamp()` for details on the interface and examples.

class brevitas.ops.autograd_ste_ops.TernarySignSteFn(*args, **kwargs)

  Bases: `Function`

  Autograd function that implements `torch.sign()` with a straight-through gradient estimator.

  `TernarySignSteFn.apply(*args)` is first aliased to `ternary_sign_ste_impl(*args)` and then wrapped by `ternary_sign_ste()` when env `BREVITAS_JIT=0`. See `ternary_sign_ste()` for details on the interface and examples.

brevitas.ops.autograd_ste_ops.abs_binary_sign_grad_impl(*args, **kwargs)

  Alias for `AbsBinarySignGradFn.apply(*args)`

brevitas.ops.autograd_ste_ops.binary_sign_ste_impl(*args, **kwargs)

  Alias for `BinarySignSteFn.apply(*args)`

brevitas.ops.autograd_ste_ops.ceil_ste_impl(*args, **kwargs)

  Alias for `CeilSteFn.apply(*args)`

brevitas.ops.autograd_ste_ops.dpu_round_ste_impl(*args, **kwargs)

  Alias for `DPURoundSteFn.apply(*args)`

brevitas.ops.autograd_ste_ops.floor_ste_impl(*args, **kwargs)

  Alias for `FloorSteFn.apply(*args)`

brevitas.ops.autograd_ste_ops.round_ste_impl(*args, **kwargs)

  Alias for `RoundSteFn.apply(*args)`

brevitas.ops.autograd_ste_ops.round_to_zero_ste_impl(*args, **kwargs)

  Alias for `RoundToZeroSteFn.apply(*args)`

brevitas.ops.autograd_ste_ops.scalar_clamp_min_ste_impl(*args, **kwargs)

  Alias for `ScalarClampMinSteFn.apply(*args)`

brevitas.ops.autograd_ste_ops.scalar_clamp_ste_impl(*args, **kwargs)

  Alias for `ScalarClampSteFn.apply(*args)`

brevitas.ops.autograd_ste_ops.tensor_clamp_ste_impl(*args, **kwargs)

  Alias for `TensorClampSteFn.apply(*args)`

brevitas.ops.autograd_ste_ops.ternary_sign_ste_impl(*args, **kwargs)

  Alias for `TernarySignSteFn.apply(*args)`