Nx

https://hexdocs.pm/nx/
Admin
Nx is a numerical computing library for Elixir that provides multi-dimensional tensors, numerical...
Tokens:488,988
Snippets:2,987
Trust Score:9
Update:1 week ago
Show doc for...
Tokens:
Raw
### Quick Start Setup Script

Source: https://hexdocs.pm/nx_quantum/readme.html

Quick start commands to set up the environment and run the example.

```bash
mise trust
mise install
mix setup
mix run examples/quantum_kernel_classifier.exs
```

--------------------------------

### Setup and run benchmark commands

Source: https://hexdocs.pm/nx_quantum/python-alternatives-benchmark-2026-03-21.html

Commands to create a virtual environment, install Python packages, and run the benchmark script for both scenarios (baseline_2q and deep_6q).

```bash
python3 -m venv .venv-bench
source .venv-bench/bin/activate
python -m pip install --upgrade pip
python -m pip install 'qiskit>=1.2,<2' 'pennylane>=0.38,<0.40' 'cirq-core>=1.3,<1.5'
python -m pip install 'autoray<0.7'
python bench/python_alternatives_benchmark.py --iterations 2000 --warmup 100 --nx-runtime-profiles cpu_portable,cpu_compiled --scenario baseline_2q
python bench/python_alternatives_benchmark.py --iterations 500 --warmup 50 --nx-runtime-profiles cpu_portable,cpu_compiled --scenario deep_6q
```

--------------------------------

### Set up the development environment with mise and mix

Source: https://hexdocs.pm/nx_quantum/contributing.html

Sets up the toolchain for the repository. Run these commands in order to trust the mise configuration, install the pinned Erlang/Elixir versions, and run the project setup tasks.

```shell
mise trust
mise install
mix setup

```

--------------------------------

### Run examples

Source: https://hexdocs.pm/nx_penalties/readme.html

Commands to run a specific example file or all examples via the provided script.

```shell
mix run examples/basic_usage.exs
./examples/run_all.sh  # Run all examples

```

--------------------------------

### new/1 examples

Source: https://hexdocs.pm/nx_penalties/NxPenalties.Pipeline.Multi.html

Examples of creating a new multi-input pipeline with and without a name option.

```elixir
pipeline = Pipeline.Multi.new()
pipeline = Pipeline.Multi.new(name: "data-aware")
```

--------------------------------

### new/1 usage examples

Source: https://hexdocs.pm/nx_penalties/NxPenalties.Pipeline.html

Examples showing how to create a new pipeline with and without a name option.

```elixir
pipeline = NxPenalties.Pipeline.new()
pipeline = NxPenalties.Pipeline.new(name: "regularization")
```

--------------------------------

### trainer/5 example

Source: https://hexdocs.pm/nx_penalties/NxPenalties.Integration.Axon.html

Example of using `NxPenalties.Integration.Axon.trainer/5` to create a training loop with an L2 penalty and run it with Axon.Loop.run.

```elixir
pipeline = NxPenalties.pipeline([
  {:l2, weight: 0.01}
])

loop = NxPenalties.Integration.Axon.trainer(
  model,
  &Axon.Losses.mean_squared_error/2,
  pipeline,
  Polaris.Optimizers.adam(learning_rate: 0.001)
)

Axon.Loop.run(loop, data, %{}, epochs: 10)
```

--------------------------------

### pipeline with options example

Source: https://hexdocs.pm/nx_penalties/NxPenalties.html

Example of creating a pipeline with configuration. The spec format allows `:weight`, `:opts`, and `:enabled` options.

```elixir
pipeline = NxPenalties.pipeline([
  {:l1, weight: 0.001},
  {:l2, weight: 0.01},
  {:entropy, weight: 0.1, opts: [mode: :penalty]}
])

```

--------------------------------

### Basic logsumexp examples

Source: https://hexdocs.pm/nx/Nx.html

Examples of Nx.logsumexp/2 with default behavior, with exp_scaling_factor as a scalar, with exp_scaling_factor as a tensor, and with a 2D tensor.

```elixir
iex> Nx.logsumexp(Nx.tensor([1, 2, 3, 4, 5, 6]))
#Nx.Tensor<
  f32
  6.4561934
>
```

```elixir
iex> Nx.logsumexp(Nx.tensor([1, 2, 3, 4, 5, 6]), exp_scaling_factor: 0.5)
#Nx.Tensor<
  f32
  5.7630463
>
```

```elixir
iex> t = Nx.tensor([1, 2, 3, 4, 5, 6])
iex> a = Nx.tensor([-1, -1, -1, 1, 1, 1])
iex> Nx.logsumexp(t, exp_scaling_factor: a)
#Nx.Tensor<
  f32
  6.356537
>
```

```elixir
iex> Nx.logsumexp(Nx.tensor([[1, 2], [3, 4], [5, 6]]))
#Nx.Tensor<
  f32
  6.4561934
>
```

--------------------------------

### compute_total/3 example with gradients

Source: https://hexdocs.pm/nx_penalties/NxPenalties.Pipeline.Multi.html

Example of using compute_total/3 inside an Nx.Defn.grad function to compute gradients.

```elixir
# Can be used with gradients
grad_fn = Nx.Defn.grad(fn inputs ->
  Pipeline.Multi.compute_total(pipeline, inputs)
end)
```

--------------------------------

### compute/3 example

Source: https://hexdocs.pm/nx_penalties/NxPenalties.Pipeline.Multi.html

Example of calling compute/3 with a pipeline and a map of input tensors, returning a tuple of total and metrics.

```elixir
{total, metrics} = Pipeline.Multi.compute(pipeline, %{
  p_logprobs: p_tensor,
  q_logprobs: q_tensor
})
```

--------------------------------

### log_penalties/3 example

Source: https://hexdocs.pm/nx_penalties/NxPenalties.Integration.Axon.html

Example showing how to add penalty metrics logging to an Axon training loop using NxPenalties.Integration.Axon.log_penalties/3.

```elixir
pipeline = NxPenalties.pipeline([
  {:l2, weight: 0.01},
  {:entropy, weight: 0.001, opts: [mode: :penalty]}
])

model
|> Axon.Loop.trainer(loss, optimizer)
|> NxPenalties.Integration.Axon.log_penalties(pipeline)
|> Axon.Loop.run(data, %{}, epochs: 10)
```

--------------------------------

### Run quick example scripts

Source: https://hexdocs.pm/nx_quantum/getting-started.html

Commands to run example scripts for quantum kernel classifier and axon hybrid train step.

```bash
mix run examples/quantum_kernel_classifier.exs
mix run examples/axon_hybrid_train_step.exs

```

--------------------------------

### Slice along axis with integer start index: Nx.slice_along_axis/3 examples

Source: https://hexdocs.pm/nx/Nx.html

Basic usage of Nx.slice_along_axis/3 with integer start indicesasi. Slices a tensor along a specified axis, optionally using named axes or a scalar tensor as start index. The examples show output tensors for axis 0 and axis 1, including named axis behavior.

```elixir
iex> Nx.slice_along_axis(Nx.iota({5, 2}), 1, 2, axis: 0)
#Nx.Tensor<
  s32[2][2]
  [
    [2, 3],
    [4, 5]
  ]
>

```

```elixir
iex> Nx.slice_along_axis(Nx.iota({2, 5}), 1, 2, axis: 1)
#Nx.Tensor<
  s32[2][2]
  [
    [1, 2],
    [6, 7]
  ]
>

```

```elixir
iex> Nx.slice_along_axis(Nx.iota({2, 5}, names: [:x, :y]), 0, 1, axis: :x)
#Nx.Tensor<
  s32[x: 1][y: 5]
  [
    [0, 1, 2, 3, 4]
  ]
>

```

```elixir
iex> Nx.slice_along_axis(Nx.iota({2, 5}, names: [:x, :y]), Nx.tensor(0), 1, axis: :x)
#Nx.Tensor<
  s32[x: 1][y: 5]
  [
    [0, 1, 2, 3, 4]
  ]
>

```

--------------------------------

### Single point example

Source: https://hexdocs.pm/nx/Nx.html

When `n: 1`, the output contains only the start value.

```elixir
iex> Nx.linspace(0, 10, n: 1)
#Nx.Tensor<
  f32[1]
  [0.0]
>
```

--------------------------------

### Nx.slice/3 basic example

Source: https://hexdocs.pm/nx/Nx.html

Slices a 1-D tensor from start index 0 with length 3, producing a 3-element tensor. The example shows the expected result for the given arguments.

```elixir
iex> Nx.slice(Nx.tensor([1, 2, 3, 4, 5, 6]), [0], [3])
#Nx.Tensor<
  s32[3]
  [1, 2, 3]
>
```

--------------------------------

### Vectorized tensor examples

Source: https://hexdocs.pm/nx/Nx.html

Vectorized tensor examples. The first vectorizes only the stop axis `:x`. The second uses different axes `:x` and `:y` for start and stop. The third uses the same axis `:x` for both.

```elixir
iex> Nx.linspace(0, Nx.vectorize(Nx.tensor([10, 20]), :x), n: 5)
#Nx.Tensor<
  vectorized[x: 2]
  f32[5]
  [
    [0.0, 2.5, 5.0, 7.5, 10.0],
    [0.0, 5.0, 10.0, 15.0, 20.0]
  ]
>
```

```elixir
iex> start = Nx.vectorize(Nx.tensor([0, 1]), :x)
iex> stop = Nx.vectorize(Nx.tensor([10, 20]), :y)
iex> Nx.linspace(start, stop, n: 5)
#Nx.Tensor<
  vectorized[x: 2][y: 2]
  f32[5]
  [
    [
      [0.0, 2.5, 5.0, 7.5, 10.0],
      [0.0, 5.0, 10.0, 15.0, 20.0]
    ],
    [
      [1.0, 3.25, 5.5, 7.75, 10.0],
      [1.0, 5.75, 10.5, 15.25, 20.0]
    ]
  ]
>
```

```elixir
iex> start = Nx.vectorize(Nx.tensor([0, 1]), :x)
iex> stop = Nx.vectorize(Nx.tensor([10, 10]), :x)
iex> Nx.linspace(start, stop, n: 5)
#Nx.Tensor<
  vectorized[x: 2]
  f32[5]
  [
    [0.0, 2.5, 5.0, 7.5, 10.0],
    [1.0, 3.25, 5.5, 7.75, 10.0]
  ]
>
```

--------------------------------

### Integer type examples

Source: https://hexdocs.pm/nx/Nx.html

Illustrates using integer types with `Nx.linspace`. The first example with `endpoint: true` shows the step calculated as (stop - start) / (n - 1), yielding a step close to 3.42 instead of 3 mountains. The second with `endpoint: false` uses a step of 3 exactly.

```elixir
iex> Nx.linspace(0, 24, n: 8, type: {:u, 8}, endpoint: true)
#Nx.Tensor<
  u8[8]
  [0, 3, 6, 10, 13, 17, 20, 24]
>
```

```elixir
iex> Nx.linspace(0, 24, n: 8, type: {:s, 32}, endpoint: false)
#Nx.Tensor<
  s32[8]
  [0, 3, 6, 9, 12, 15, 18, 21]
>
```

--------------------------------

### Verify build and run tests (make mac, mix test)

Source: https://hexdocs.pm/nx_tflite_mob/build_mac_tflite.html

Build the host-side NIF and run tests. `make mac` produces priv/mac/libtflite_nif.so, then copy it to priv/native/ and run `mix test` (16 tests, 0 failures).

```bash
cd ~/code/nx_tflite_mob
make mac          # produces priv/mac/libtflite_nif.so (~50 KB)
cp priv/mac/libtflite_nif.so priv/native/libtflite_nif.so
mix test          # 16 tests, 0 failures

```

--------------------------------

### Nx.new_axis/3 examples for positive and negative axes

Source: https://hexdocs.pm/nx/Nx.html

Examples for Nx.new_axis/3 that add a new axis of size 1 at a specified position with an optional name. The first block shows adding at axes 0, 1, and 2. The second block demonstrates that a negative axis index starts from the back, here -1 inserting at the end. Tensor shapes and contents are displayed as #Nx.Tensor<...> output.

```elixir
iex> t = Nx.tensor([[1, 2, 3], [4, 5, 6]])
iex> Nx.new_axis(t, 0, :new)
#Nx.Tensor<
  s32[new: 1][2][3]
  [
    [
      [1, 2, 3],
      [4, 5, 6]
    ]
  ]
>
iex> Nx.new_axis(t, 1, :new)
#Nx.Tensor<
  s32[2][new: 1][3]
  [
    [
      [1, 2, 3]
    ],
    [
      [4, 5, 6]
    ]
  ]
>
iex> Nx.new_axis(t, 2, :new)
#Nx.Tensor<
  s32[2][3][new: 1]
  [
    [
      [1],
      [2],
      [3]
    ],
    [
      [4],
      [5],
      [6]
    ]
  ]
>
```

```elixir
iex> t = Nx.tensor([[1, 2, 3], [4, 5, 6]])
iex> Nx.new_axis(t, -1, :new)
#Nx.Tensor<
  s32[2][3][new: 1]
  [
    [
      [1],
      [2],
      [3]
    ],
    [
      [4],
      [5],
      [6]
    ]
  ]
>
```

--------------------------------

### build_train_step/4 example

Source: https://hexdocs.pm/nx_penalties/NxPenalties.Integration.Axon.html

Example of building a training step with penalties. Creates a pipeline with L2 and entropy penalties, builds the step, initializes state, and runs a training step.

```elixir
pipeline = NxPenalties.pipeline([
  {:l2, weight: 0.01},
  {:entropy, weight: 0.001, opts: [mode: :penalty]}
])

{init_fn, step_fn} = NxPenalties.Integration.Axon.build_train_step(
  model,
  &Axon.Losses.mean_squared_error/2,
  pipeline,
  Polaris.Optimizers.adam(learning_rate: 0.001)
)

# Initialize state
state = init_fn.(Nx.template({1, 10}, :f32), %{})

# Run training step
{new_state, metrics} = step_fn.(state, {x_batch, y_batch})
```

--------------------------------

### Nx.floor: compute floor of tensor elements

Source: https://hexdocs.pm/nx/Nx.html

We have two examples of the floor operation. The first shows that non-floating tensors are returned as is if given. The second shows that floating tensors are floored towards negative infinity and the type is kept. Use the floor function to get the integer part of float elements.

```elixir
iex> Nx.floor(Nx.tensor([-1, 0, 1], names: [:x]))
#Nx.Tensor<
  s32[x: 3]
  [-1, 0, 1]
>

```

```elixir
iex> Nx.floor(Nx.tensor([-1.5, -0.5, 0.5, 1.5], names: [:x]))
#Nx.Tensor<
  f32[x: 4]
  [-2.0, -1.0, 0.0, 1.0]
>

```

--------------------------------

### Build and compute a pipeline with NxPenalties.Pipeline

Source: https://hexdocs.pm/nx_penalties/NxPenalties.Pipeline.html

Shows how to create a pipeline, add three penalties (l1, l2, entropy) with weights and options, and compute the total penalty and metrics. The entropy penalty uses opts [mode: :penalty].

```elixir
pipeline =
  NxPenalties.Pipeline.new()
  |> NxPenalties.Pipeline.add(:l1, &NxPenalties.Penalties.l1/2, weight: 0.001)
  |> NxPenalties.Pipeline.add(:l2, &NxPenalties.Penalties.l2/2, weight: 0.01)
  |> NxPenalties.Pipeline.add(:entropy, &NxPenalties.Divergences.entropy/2,
       weight: 0.1, opts: [mode: :penalty])

{total_penalty, metrics} = NxPenalties.Pipeline.compute(pipeline, tensor)
```

--------------------------------

### YOLOv8n TFLite model export and COCO labels setup

Source: https://hexdocs.pm/nx_tflite_mob/yolo_walkthrough.html

Set up the YOLOv8n model and COCO labels. The commands install Ultralytics, export the model (producing several TFLite variants), copy the needed .tflite files into the app's priv/yolo directory, and define the 80 COCO class labels in Elixir syntax.

```bash
pip install ultralytics
yolo export model=yolov8n.pt format=tflite int8=True imgsz=640
# → yolov8n_saved_model/yolov8n_full_integer_quant.tflite
# → yolov8n_saved_model/yolov8n_float16.tflite
# → yolov8n_saved_model/yolov8n_float32.tflite

```

```bash
mkdir -p priv/yolo
cp yolov8n_saved_model/yolov8n_full_integer_quant.tflite priv/yolo/
cp yolov8n_saved_model/yolov8n_float16.tflite priv/yolo/

```

```elixir
@coco ~w(
  person bicycle car motorcycle airplane bus train truck boat
  traffic-light fire-hydrant stop-sign parking-meter bench bird cat
  dog horse sheep cow elephant bear zebra giraffe backpack umbrella
  handbag tie suitcase frisbee skis snowboard sports-ball kite
  baseball-bat baseball-glove skateboard surfboard tennis-racket
  bottle wine-glass cup fork knife spoon bowl banana apple sandwich
  orange broccoli carrot hot-dog pizza donut cake chair couch
  potted-plant bed dining-table toilet tv laptop mouse remote
  keyboard cell-phone microwave oven toaster sink refrigerator book
  clock vase scissors teddy-bear hair-drier toothbrush
)

```

--------------------------------

### schedule_weights example with Axon.Loop

Source: https://hexdocs.pm/nx_penalties/NxPenalties.Integration.Axon.html

Shows how to use schedule_weights with a weight schedule and Axon.Loop, including storing the pipeline in state on the first epoch.

```elixir
schedule_fn = NxPenalties.Integration.Axon.weight_schedule(%{
  l2: {:linear, 0.0, 0.01},
  kl: {:warmup, 0.1, 10}
})

model
|> Axon.Loop.trainer(loss, optimizer)
|> Axon.Loop.handle_event(:epoch_started, fn state ->
  # Store pipeline in state on first epoch
  pipeline = state[:handler_metadata][:pipeline] || initial_pipeline
  {:continue, put_in(state[:handler_metadata][:pipeline], pipeline)}
end)
|> NxPenalties.Integration.Axon.schedule_weights(:pipeline, schedule_fn, total_epochs: 50)
|> Axon.Loop.run(data, %{}, epochs: 50)
```

--------------------------------

### total_grad_norm example

Source: https://hexdocs.pm/nx_penalties/NxPenalties.GradientTracker.html

Example of using `total_grad_norm/2` with a pipeline and tensor.

```elixir
norm = GradientTracker.total_grad_norm(pipeline, tensor)
```

--------------------------------

### Nx.phase examples

Source: https://hexdocs.pm/nx/Nx.html

Examples of calling Nx.phase/1 on a complex number, a real number, and a vector literal. The last example uses the ~VEC sigil to create a complex vector.

```elixir
iex> Nx.phase(Complex.new(1, 2))
#Nx.Tensor<
  f32
  1.1071488
>

```

```elixir
iex> Nx.phase(1)
#Nx.Tensor<
  f32
  0.0
>

```

```elixir
iex> import Nx, only: [sigil_VEC: 2]
iex> Nx.phase(~VEC[1+2i -2+1i])
#Nx.Tensor<
  f32[2]
  [1.1071488, 2.6779451]
>

```

--------------------------------

### weight_schedule/1 example

Source: https://hexdocs.pm/nx_penalties/NxPenalties.Integration.Axon.html

Example of creating a weight schedule function with linear, warmup, and constant schedules, and using it to compute weights for a given epoch.

```elixir
schedule_fn = NxPenalties.Integration.Axon.weight_schedule(%{
  l2: {:linear, 0.0, 0.01},        # Ramp L2 from 0 to 0.01
  kl: {:warmup, 0.1, 10},          # Warm up KL over 10 epochs
  entropy: {:constant, 0.001}      # Keep entropy constant
})

# Use in training loop
epoch = 5
total_epochs = 100
weights = schedule_fn.(epoch, total_epochs)
# => %{l2: 0.0005, kl: 0.05, entropy: 0.001}
```

--------------------------------

### remove/2 example

Source: https://hexdocs.pm/nx_penalties/NxPenalties.Pipeline.Multi.html

Example of removing a penalty named :kl from the pipeline.

```elixir
pipeline = Pipeline.Multi.remove(pipeline, :kl)
```

--------------------------------

### Set default backend via Mix.install config

Source: https://hexdocs.pm/nx/Nx.html

Install Nx and set EXLA.Backend as the default backend inside a Mix.install/2 call in a notebook or script. The config option passes the backend configuration to Nx at startup. Ensure EXLA is added as a dependency (here only Nx is listed).

```elixir
Mix.install(
  [
    {:nx, ">= 0.0.0"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)

```

--------------------------------

### js_divergence function example

Source: https://hexdocs.pm/nx_penalties/NxPenalties.Divergences.html

Example of using the js_divergence function with log-probabilities.

```elixir
iex> p_logprobs = Nx.log(Nx.tensor([0.4, 0.3, 0.2, 0.1]))
iex> q_logprobs = Nx.log(Nx.tensor([0.25, 0.25, 0.25, 0.25]))
iex> NxPenalties.Divergences.js_divergence(p_logprobs, q_logprobs)
```

--------------------------------

### Nx.window_max examples with strides, padding, and window_dilations

Source: https://hexdocs.pm/nx/Nx.html

Basic examples from the function's documentation that compute windowed maximums on tensors, showing output with integer and float typesasi as well as default, explicit, and valid padding configurations. The examples are self-contained and demonstrate different argument combinations for strides and padding.

```elixir
iex> Nx.window_max(Nx.tensor([[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]]), {1, 2, 1})
#Nx.Tensor<
  s32[2][1][3]
  [
    [
      [4, 5, 6]
    ],
    [
      [4, 5, 6]
    ]
  ]
>
```

```elixir
iex> t = Nx.tensor([[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]])
iex> Nx.window_max(t, {2, 2, 1}, strides: [1, 2, 3], padding: [{0, 1}, {2, 0}, {1, 1}])
#Nx.Tensor<
  s32[2][2][2]
  [
    [
      [-2147483648, -2147483648],
      [-2147483648, 6]
    ],
    [
      [-2147483648, -2147483648],
      [-2147483648, 6]
    ]
  ]
>
```

```elixir
iex> t = Nx.tensor([[[4.0, 2.0, 3.0], [2.0, 5.0, 6.5]], [[1.2, 2.2, 3.2], [4.0, 5.0, 6.2]]])
iex> Nx.window_max(t, {2, 1, 1}, strides: [2, 1, 1], padding: [{1, 1}, {0, 0}, {1, 1}])
#Nx.Tensor<
  f32[2][2][5]
  [
    [
      [-Inf, 4.0, 2.0, 3.0, -Inf],
      [-Inf, 2.0, 5.0, 6.5, -Inf]
    ],
    [
      [-Inf, 1.2, 2.2, 3.2, -Inf],
      [-Inf, 4.0, 5.0, 6.2, -Inf]
    ]
  ]
>
```

```elixir
iex> t = Nx.tensor([[[4, 2, 1, 3], [4, 2, 1, 7]], [[1, 2, 5, 7], [1, 8, 9, 2]]])
iex> opts = [strides: [2, 1, 1], padding: :valid, window_dilations: [1, 2, 2]]
iex> Nx.window_max(t, {1, 1, 2}, opts)
#Nx.Tensor<
  s32[1][2][2]
  [
    [
      [4, 3],
      [4, 7]
    ]
  ]
>
```

--------------------------------

### put_slice/3 with dynamic start indices

Source: https://hexdocs.pm/nx/Nx.html

Dynamic start indices: each start index can be a scalar tensor. The slice tensor is f32, which changes the result type. Works like slice/3.

```elixir
iex> t = Nx.tensor([[1, 2, 3], [4, 5, 6]])
iex> Nx.put_slice(t, [Nx.tensor(0), Nx.tensor(1)], Nx.tensor([[10.0, 11.0]]))
#Nx.Tensor<
  f32[2][3]
  [
    [1.0, 10.0, 11.0],
    [4.0, 5.0, 6.0]
  ]
>
```

--------------------------------

### Create and compute a multi-input pipeline with Pipeline.Multi

Source: https://hexdocs.pm/nx_penalties/NxPenalties.Pipeline.Multi.html

Shows how to build a multi-input pipeline by adding penalties with named inputs and weights, then compute the total penalty and metrics using a map of named inputs. The example uses KL divergence and consistency penalties.

```elixir
# Create multi-input pipeline
pipeline =
  Pipeline.Multi.new()
  |> Pipeline.Multi.add(:kl, &Divergences.kl_divergence/3,
    inputs: [:student_logprobs, :teacher_logprobs],
    weight: 0.1
  )
  |> Pipeline.Multi.add(:consistency, &Constraints.consistency/3,
    inputs: [:clean_output, :noisy_output],
    weight: 0.2,
    opts: [metric: :mse]
  )

# Compute with named inputs
{total, metrics} = Pipeline.Multi.compute(pipeline, %{
  student_logprobs: student_out,
  teacher_logprobs: teacher_out,
  clean_output: clean,
  noisy_output: noisy
})
```

--------------------------------

### set_enabled/3 usage example

Source: https://hexdocs.pm/nx_penalties/NxPenalties.Pipeline.html

Example showing how to disable a penalty named :l1.

```elixir
pipeline = Pipeline.set_enabled(pipeline, :l1, false)
```

--------------------------------

### Nx.window_scatter_min basic example

Source: https://hexdocs.pm/nx/Nx.html

Shows two runs of Nx.window_scatter_min with different strides and window sizes, reproducing the documented outputs. The output tensor is initialized with the given init_value (0) and overlapped indices accumulate source values.

```elixir
iex> t = Nx.tensor([
...>   [7, 2, 5, 3, 10, 2],
...>   [3, 8, 9, 3, 4, 2],
...>   [1, 5, 7, 5, 6, 1],
...>   [0, 6, 2, 7, 2, 8]
...> ])
iex> opts = [strides: [2, 3], padding: :valid]
iex> Nx.window_scatter_min(t, Nx.tensor([[2, 6], [3, 1]]), 0, {2, 3}, opts)
#Nx.Tensor<
  s32[4][6]
  [
    [0, 2, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 6],
    [0, 0, 0, 0, 0, 1],
    [3, 0, 0, 0, 0, 0]
  ]
>
```

```elixir
iex> t = Nx.tensor([
...>   [7, 2, 5, 3, 8],
...>   [3, 8, 9, 3, 4],
...>   [1, 5, 7, 5, 6],
...>   [0, 6, 2, 10, 2]
...> ])
iex> opts = [strides: [2, 2], padding: :valid]
iex> Nx.window_scatter_min(t, Nx.tensor([[2, 6], [3, 1]]), 0, {2, 3}, opts)
#Nx.Tensor<
  s32[4][5]
  [
    [0, 2, 0, 0, 0],
    [0, 0, 0, 6, 0],
    [0, 0, 0, 0, 0],
    [3, 0, 0, 0, 1]
  ]
>
```

--------------------------------

### set_enabled/3 example

Source: https://hexdocs.pm/nx_penalties/NxPenalties.Pipeline.Multi.html

Example usage of set_enabled/3 to disable the penalty named :kl.

```elixir
pipeline = Pipeline.Multi.set_enabled(pipeline, :kl, false)
```

--------------------------------

### Nx.window_mean examples with strides, padding, and dilations

Source: https://hexdocs.pm/nx/Nx.html

Examples of `Nx.window_mean/3` with default options (padding :valid, dilations 1) and with custom `:strides`, `:padding` (including padding with zeros) and `:window_dilations`. Padding values must be non-negative; negative pads are not allowed. The output shape is determined by window, strides, dilations, and padding; :same pads so output shape equals input shape.

```elixir
iex> t = Nx.tensor([[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]])
iex> Nx.window_mean(t, {1, 2, 1})
#Nx.Tensor<
  f32[2][1][3]
  [
    [
      [2.5, 3.5, 4.5]
    ],
    [
      [2.5, 3.5, 4.5]
    ]
  ]
>
```

```elixir
iex> t = Nx.tensor([[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]])
iex> Nx.window_mean(t, {2, 2, 1}, strides: [1, 2, 3], padding: [{0, 1}, {2, 0}, {1, 1}])
#Nx.Tensor<
  f32[2][2][2]
  [
    [
      [0.0, 0.0],
      [0.0, 4.5]
    ],
    [
      [0.0, 0.0],
      [0.0, 2.25]
    ]
  ]
>
```

```elixir
iex> t = Nx.tensor([[[4.0, 2.0, 3.0], [2.0, 5.0, 6.5]], [[1.2, 2.2, 3.2], [4.0, 5.0, 6.2]]])
iex> Nx.window_mean(t, {2, 1, 1}, strides: [2, 1, 1], padding: [{1, 1}, {0, 0}, {1, 1}])
#Nx.Tensor<
  f32[2][2][5]
  [
    [
      [0.0, 2.0, 1.0, 1.5, 0.0],
      [0.0, 1.0, 2.5, 3.25, 0.0]
    ],
    [
      [0.0, 0.6, 1.1, 1.6, 0.0],
      [0.0, 2.0, 2.5, 3.1, 0.0]
    ]
  ]
>
```

```elixir
iex> t = Nx.tensor([[[4, 2, 1, 3], [4, 2, 1, 7]], [[1, 2, 5, 7], [1, 8, 9, 2]]])
iex> opts = [strides: [2, 1, 1], padding: :valid, window_dilations: [1, 2, 1]]
iex> Nx.window_mean(t, {1, 1, 2}, opts)
#Nx.Tensor<
  f32[1][2][3]
  [
    [
      [3.0, 1.5, 2.0],
      [3.0, 1.5, 4.0]
    ]
  ]
>
```

```elixir
iex> t = Nx.tensor([[[4, 2, 1, 3], [4, 2, 1, 7]], [[1, 2, 5, 7], [1, 8, 9, 2]]])
iex> opts = [strides: [2, 1, 1], padding: :valid, window_dilations: [1, 2, 2]]
iex> Nx.window_mean(t, {1, 1, 2}, opts)
#Nx.Tensor<
  f32[1][2][2]
  [
    [
      [2.5, 2.5],
      [2.5, 4.5]
    ]
  ]
>
```

--------------------------------

### output_magnitude_penalty/2 example

Source: https://hexdocs.pm/nx_penalties/NxPenalties.GradientPenalty.html

Example usage of output_magnitude_penalty, penalizing output magnitude with a target of 1.0.

```elixir
penalty = GradientPenalty.output_magnitude_penalty(model_output, target: 1.0)
```

--------------------------------

### Nx.window_product basic usage, strides, padding, and dilation

Source: https://hexdocs.pm/nx/Nx.html

Basic window_product examples showing the function applied to integer tensors with default optionsanding without specifying strides or padding. The first call uses a window of {1, 2, 1} and returns s32[2][1][3] output. The second call combines strides: [1, 2, 3] and padding: [{0, 1}, {2, 0}, {1, 1}] to produce s32[2][2][2]. The third example uses float tensors with strides and padding producing f32[2][2][5]. The fourth example demonstrates the use of explicit options with window_dilations and padding: :valid producing s32[1][2][2]. Each example includes the exact resultant tensor.

```elixir
iex> Nx.window_product(Nx.tensor([[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]]), {1, 2, 1})
#Nx.Tensor<
  s32[2][1][3]
  [
    [
      [4, 10, 18]
    ],
    [
      [4, 10, 18]
    ]
  ]
>
```

```elixir
iex> t = Nx.tensor([[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]])
iex> Nx.window_product(t, {2, 2, 1}, strides: [1, 2, 3], padding: [{0, 1}, {2, 0}, {1, 1}])
#Nx.Tensor<
  s32[2][2][2]
  [
    [
      [1, 1],
      [1, 324]
    ],
    [
      [1, 1],
      [1, 18]
    ]
  ]
>
```

```elixir
iex> t = Nx.tensor([[[4.0, 2.0, 3.0], [2.0, 5.0, 6.5]], [[1.2, 2.2, 3.2], [4.0, 5.0, 6.2]]])
iex> Nx.window_product(t, {2, 1, 1}, strides: [2, 1, 1], padding: [{1, 1}, {0, 0}, {1, 1}])
#Nx.Tensor<
  f32[2][2][5]
  [
    [
      [1.0, 4.0, 2.0, 3.0, 1.0],
      [1.0, 2.0, 5.0, 6.5, 1.0]
    ],
    [
      [1.0, 1.2, 2.2, 3.2, 1.0],
      [1.0, 4.0, 5.0, 6.2, 1.0]
    ]
  ]
>
```

```elixir
iex> t = Nx.tensor([[[4, 2, 1, 3], [4, 2, 1, 7]], [[1, 2, 5, 7], [1, 8, 9, 2]]])
iex> opts = [strides: [2, 1, 1], padding: :valid, window_dilations: [1, 2, 2]]
iex> Nx.window_product(t, {1, 1, 2}, opts)
#Nx.Tensor<
  s32[1][2][2]
  [
    [
      [4, 6],
      [4, 14]
    ]
  ]
>
```

--------------------------------

### Create and compute a multi-input pipeline with NxPenalties.Pipeline.Multi

Source: https://hexdocs.pm/nx_penalties/readme.html

Shows how to build a multi-input pipeline using Multi.new/0, add penalties with named inputs and weights, then compute the total penalty and per-metric metrics with Multi.compute/2. Also includes the gradient-safe variant Multi.compute_total/2.

```elixir
alias NxPenalties.Pipeline.Multi

# Create multi-input pipeline
pipeline =
  Multi.new()
  |> Multi.add(:kl, &NxPenalties.Divergences.kl_divergence/3,
      inputs: [:student_logprobs, :teacher_logprobs],
      weight: 0.1
    )
  |> Multi.add(:consistency, &NxPenalties.Constraints.consistency/3,
      inputs: [:clean_output, :noisy_output],
      weight: 0.2,
      opts: [metric: :mse]
    )

# Compute with named inputs
{total, metrics} = Multi.compute(pipeline, %{
  student_logprobs: student_out,
  teacher_logprobs: teacher_out,
  clean_output: clean,
  noisy_output: noisy
})

# Gradient-safe version
total = Multi.compute_total(pipeline, inputs_map)
```

--------------------------------

### gradient_penalty/3 example

Source: https://hexdocs.pm/nx_penalties/NxPenalties.GradientPenalty.html

Example usage of gradient_penalty, penalizing gradients that deviate from norm 1.

```elixir
# Penalize gradients that deviate from norm 1
penalty = GradientPenalty.gradient_penalty(
  fn x -> Nx.sum(x) end,
  tensor,
  target_norm: 1.0
)
```
Tokens:
Raw