AD of a model with respect to an attribute of a class

[AlvaroDFC]

Hello, I have recently started to use Jax and I wanted to apply it to an automatic differentiation problem. I have a function that takes three inputs: 2 classes and 1 scalar. I want to take the sensitivity of my output (6x6 matrix) with respect to an attribute (scalar) of one of my classes. At the moment to work with partial derivatives, I am passing a fourth argument, Cf (Cf which is a tensor, and the line that calculates it is commented out in MT method). The goal is to remove the input Cf and calculate the derivative with respect to Fiber.EL. Here is an example code:

from Material_class import Mat_Database
from jax import numpy as jnp
from jax import *
config.update("jax_enable_x64", True)

def MT(Fiber,Matrix,Vf,Cf):
    if Vf > 0.99999999:
        Vf = 0.99999999

    def eshelbyCil(Num):
        P = jnp.zeros((6, 6))

        P = P.at[1, 0].set(Num / (2. * (1. - Num)))
        P = P.at[2, 0].set(Num / (2. * (1. - Num)))
        P = P.at[1, 1].set((1. / (2. * (1. - Num))) * (3. / 4. + (1. - 2. * Num) / 2.))
        P = P.at[2, 2].set(P[1, 1])
        P = P.at[1, 2].set((1. / (2. * (1. - Num))) * (1. / 4. - (1. - 2. * Num) / 2.))
        P = P.at[2, 1].set(P[1, 2])
        # -- Engineering shear strains - multiply tensorial P by 2
        P = P.at[3, 3].set((2. / (2. * (1. - Num))) * (1. / 4. + (1. - 2. * Num) / 2.))
        P = P.at[4, 4].set(2. / 4.)
        P = P.at[5, 5].set(2. / 4.)

        return P

    # Stiffness tensors of constituents. Returns a 6x6 matrix that depends on the material properties
    # Cf = Fiber.Stiffness()
    Cm = Matrix.Stiffness()

    # Eshelby stiffness tensor, eq 3.41 of Arnold 2020
    C = jnp.linalg.solve(Cm, (Cm - Cf))
    # Eshelby Tensor for cylindrical fiber
    P = eshelbyCil(Matrix.nuL)

    # Strain concentration tensor of Eshelby inclusion
    Af_dilute = jnp.linalg.inv(jnp.eye(6) - jnp.matmul(P , C))
    # Strain concentration tensor of MT theory
    Af_MT = jnp.matmul(Af_dilute, jnp.linalg.inv((1 - Vf) * jnp.eye(6) + Vf * Af_dilute))

    # Homogenized stiffness
    Cstar = Cm + Vf * jnp.matmul((Cf - Cm) , Af_MT)

    return Cstar

Fiber   = Mat_Database.SandV_carbon() # Returns material properties (EL, ET, GL, GT, nuL,nuT)
Matrix  = Mat_Database.SandV_epoxy() # Returns material properties (EL, ET, GL, GT, nuL,nuT)
Vf      = 0.6

Cf = Fiber.Stiffness()

Cstar = MT(Fiber,Matrix,Vf,Cf)

jnp.set_printoptions(precision=4,linewidth=150)
print(Cstar)

# Testing AD with MT ##
# Test with jacfwd

f = lambda Cf: MT(Fiber,Matrix,Vf,Cf)
J = jacfwd(f)(Cf)

print("jacfwd result, with shape", J.shape)
print(J)

[jakevdp]

The easiest way to do this would be to ensure your class is registered as a pytree, and then define a function with respect to its flattened contents and take the gradient with respect to the desired argument. That's admittedly not a very clean solution here; #10614 is a feature request for more expressive gradient argument specifications that might allow a cleaner solution in this case.

[AlvaroDFC]

Thank you for the answer. I have simplified my example so that it contains all the operations I need and it can directly be reproduced in anyone's machine.

• I have registered my class 'Material' as a pytree.
• I am not sure I have defined the function with respect to its flattened content. Does that mean that I cannot pass my class as input arguments, but rather I need to rewrite my method to only take the flattened content?
• I have tried the following, but does not work (and all the combinations i could think of):

from jax import numpy as jnp
from jax.tree_util import register_pytree_node,tree_unflatten, tree_flatten
import jax
jax.config.update("jax_enable_x64", True)

def MT(Fiber,Matrix,Vf):
    # Stiffness tensors of constituents
    Cf = Fiber.Stiffness()
    Cm = Matrix.Stiffness()

    # Homogenized stiffness
    Cstar = Cm*(1-Vf) + Vf * Cf

    return Cstar

class Material:
    
    def __init__(self, name=[],EL=0, ET=0):

        self.name   = name
        # L stands for Longitudinal == x1 - direction
        # T stands for Tranversal == x2,x3 - direction
        self.EL     = EL    # Young Modulus
        self.ET     = ET    # Young Modulus

    def Stiffness(self):

        S = jnp.zeros((3,3),dtype='float64')
        S = S.at[0,0].set(1./self.EL)
        S = S.at[1,1].set(1./self.ET)
        S = S.at[2,2].set(1.)

        b = jnp.eye(3)
        C = jnp.zeros((3,3))
        
        for i in range(len(b)):
            aux = jnp.linalg.solve(S,b[:,i])
            C = C.at[:,i].set(aux)

        return C


def flatten_counter(tree):
    children = (tree.EL, tree.ET, tree.Stiffness())
    aux_data = None
    return (children, aux_data)

def unflatten_counter(data,children):
    return Material(*children, data)


##########################
Fiber   = Material('Fiber', 100., 10.)
Matrix  = Material('Matrix', 10., 10.)
Vf      = 0.5

register_pytree_node(Material, flatten_counter, unflatten_counter)
children,data   = flatten_counter(Fiber)
# tree            = unflatten_counter(data,children)

# Flatten the custom object
my_leavesF, my_treedefF = tree_flatten(Fiber)
my_leavesM, my_treedefM = tree_flatten(Matrix)

print(my_leavesF)
print(my_treedefF)

def activate(x): # Dummy function
    return jnp.sin(x)

# Test example
grad2 = jax.grad(activate)(my_leavesF[0]) # This works
print(grad2)


# Real example
f = lambda Fiber: MT(Fiber,Matrix,Vf)
J = jax.grad(f)(my_leavesF[0]) # AttributeError: JVPTracer has no attribute Stiffness

I tried to understand issue #10614, but my knowledge of jax is definitely not enough. Thank you in advance.

[jakevdp]

What I mean by "define a function with respect to its flattened contents" is defining a function that accepts arrays or scalars as arguments. For example:

def f(EL, ET):
  fiber = Material("Fiber", EL, ET)
  return fiber.Stiffness().sum()

print(jax.grad(f, argnums=1)(1.0, 1.0))

It's not very convenient, but since there's currently no way to construct the gradient with respect to a class attribute, this is the best workaround.

[AlvaroDFC]

This is great! Thank you for your help.