custom_grad and vectorized inputs

Arn3n · January 28, 2024, 2:50pm

Hello,

I’ve been using the NX library, and have been trying to write a custom_grad function that works off of a vectorized input tensor. I’m having trouble getting the function to pass both during autograd phase AND during runtime, usually due to shape issues. I’m not looking for someone to debug my code – instead, I’m wondering if there are any non-trivial examples of custom_grad that I could look at that work off of vectorized tensors.

polvalente · January 28, 2024, 3:10pm

github.com

elixir-nx/nx/blob/main/nx/lib/nx/lin_alg/svd.ex

defmodule Nx.LinAlg.SVD do
  @moduledoc false

  # Implementation of Singular Value Decompositon
  # Inspired by the Jax code in https://github.com/google/jax/blob/ba557d5e1beb480851117a003ebf76c0ed2249e0/jax/_src/lax/svd.py
  # QDWH is QR-based Dynamically Weighted Halley iteration
  #
  # References (same as Jax, above):
  #
  # * Nakatsukasa, Yuji, and Nicholas J. Higham.
  #   "Stable and efficient spectral divide and conquer algorithms for the symmetric
  #   eigenvalue decomposition and the SVD." SIAM Journal on Scientific Computing 35,
  #   no. 3 (2013): A1325-A1349.
  #   https://epubs.siam.org/doi/abs/10.1137/120876605
  #
  # * Nakatsukasa, Yuji, Zhaojun Bai, and François Gygi.
  #   "Optimizing Halley's iteration for computing the matrix polar decomposition."
  #   SIAM Journal on Matrix Analysis and Applications 31, no. 5 (2010): 2700-2720.
  #   https://epubs.siam.org/doi/abs/10.1137/090774999

This file has been truncated. show original

This one is the example that came to mind for me, but the input is explicitly not vectorized there. Not sure if it helps