<jupyter_start><jupyter_text>In this notebook, we'll look at Tensorflow's tools for computing a jacobian in the context of an SR algorithm.We'll initialize some input, a network, and try to compute the jacobian in several different ways.<jupyter_code>import tensorflow as tf
import time, timeit

N_WALKERS = 100
DIM = 3
N_PARTICLES = 4
x_input = tf.random.uniform(shape=(N_WALKERS, N_PARTICLES, DIM))
class DeepSetsWavefunction(tf.keras.models.Model):
    """Create a neural network eave function in N dimensions

    Boundary condition, if not supplied, is gaussian in every dimension

    Extends:
        tf.keras.models.Model
    """
    def __init__(self, ndim : int, nparticles: int, mean_subtract : bool, boundary_condition :tf.keras.layers.Layer = None):
        '''Deep Sets wavefunction for symmetric particle wavefunctions

        Implements a deep set network for multiple particles in the same system

        Arguments:
            ndim {int} -- Number of dimensions
            nparticles {int} -- Number of particls

        Keyword Arguments:
            boundary_condition {tf.keras.layers.Layer} -- [description] (default: {None})

        Raises:
            Exception -- [description]
        '''
        tf.keras.models.Model.__init__(self)

        self.ndim = ndim
        if self.ndim < 1 or self.ndim > 3:
           raise Exception("Dimension must be 1, 2, or 3 for DeepSetsWavefunction")

        self.nparticles = nparticles

        self.mean_subtract = mean_subtract


        n_filters_per_layer = 8
        n_layers            = 1
        bias                = True
        activation          = tf.keras.activations.tanh


        self.individual_net = tf.keras.models.Sequential()

        self.individual_net.add(
            tf.keras.layers.Dense(n_filters_per_layer,
                use_bias = bias)
            )

        for l in range(n_layers):
            self.individual_net.add(
                tf.keras.layers.Dense(n_filters_per_layer,
                    use_bias    = bias,
                    activation = activation)
                )


        self.aggregate_net = tf.keras.models.Sequential()

        for l in range(n_layers):
            self.individual_net.add(
                tf.keras.layers.Dense(n_filters_per_layer,
                    use_bias    = bias,
                    activation = activation)
                )
        self.aggregate_net.add(tf.keras.layers.Dense(1,
            use_bias = False))


    @tf.function(experimental_compile=False)
    def call(self, inputs, trainable=None):
        # Mean subtract for all particles:
        if self.nparticles > 1 and self.mean_subtract:
            mean = tf.reduce_mean(inputs, axis=1)
            xinputs = inputs - mean[:,None,:]
        else:
            xinputs = inputs

        x = []
        for p in range(self.nparticles):
            x.append(self.individual_net(xinputs[:,p,:]))

        x = tf.add_n(x)
        x = self.aggregate_net(x)

        # Compute the initial boundary condition, which the network will slowly overcome
        # boundary_condition = tf.math.abs(self.normalization_weight * tf.reduce_sum(xinputs**self.normalization_exponent, axis=(1,2))
        boundary_condition = -1. * tf.reduce_sum(xinputs**2, axis=(1,2))
        boundary_condition = tf.reshape(boundary_condition, [-1,1])


        return x + boundary_condition

    def n_parameters(self):
        return tf.reduce_sum( [ tf.reduce_prod(p.shape) for p in self.trainable_variables ])

wavefunction = DeepSetsWavefunction(ndim=DIM, nparticles=N_PARTICLES, mean_subtract=True)
output = wavefunction(x_input)<jupyter_output><empty_output><jupyter_text>Now, we have a compiled wavfunction with a number of parameters<jupyter_code>print(wavefunction.summary())
# In general, we can compute the gradient with respect to the inputs:

with tf.GradientTape() as tape:
    log_psiw = wavefunction(x_input)

# By default, this essentially SUMS over the dimension of log_psiw
print(log_psiw.shape)
grads = tape.gradient(log_psiw, wavefunction.trainable_variables)
for g in grads:
    print(g.shape)<jupyter_output>(3, 8)
(8,)
(8, 8)
(8,)
(8, 8)
(8,)
(8, 1)<jupyter_text>We can also compute the jacobian, instead of the gradient, which in this case is the gradient but for only one walker at a time:<jupyter_code>with tf.GradientTape() as tape:
    log_psiw = wavefunction(x_input)

# By default, this essentially SUMS over the dimension of log_psiw
print(log_psiw.shape)
jac = tape.jacobian(log_psiw, wavefunction.trainable_variables)
for j in jac:
    print(j.shape)<jupyter_output>(100, 1, 3, 8)
(100, 1, 8)
(100, 1, 8, 8)
(100, 1, 8)
(100, 1, 8, 8)
(100, 1, 8)
(100, 1, 8, 1)<jupyter_text>We can verify this, when average, matches the total gradient:<jupyter_code>for j, g in zip(jac, grads):
#     print(j)
#     print(g)
#     print(tf.reduce_sum(j, axis=(0,1)))
    diff = tf.abs(g - tf.reduce_sum(j, axis=(0,1)))
    print(tf.reduce_max(diff))<jupyter_output>tf.Tensor(8.1956387e-07, shape=(), dtype=float32)
tf.Tensor(3.0517578e-05, shape=(), dtype=float32)
tf.Tensor(8.6426735e-07, shape=(), dtype=float32)
tf.Tensor(4.5776367e-05, shape=(), dtype=float32)
tf.Tensor(8.34465e-07, shape=(), dtype=float32)
tf.Tensor(7.6293945e-05, shape=(), dtype=float32)
tf.Tensor(0.0, shape=(), dtype=float32)<jupyter_text>For every layer, the jacobian in this case is quite well matching the gradients if you sum over the input dimension. It seems like we should be able to compute this per walker too:<jupyter_code>with tf.GradientTape(persistent=True) as tape:
    log_psiw = wavefunction(x_input)

    split = tf.split(log_psiw, (1, N_WALKERS-1))
    
# print(split)
# By default, this essentially SUMS over the dimension of log_psiw
grad = tape.gradient(split[1], wavefunction.trainable_variables)

@tf.function
def jacobian_comp(inputs, _wavefunction):

    with tf.GradientTape() as tape:
        log_psiw = _wavefunction(inputs)

    # By default, this essentially SUMS over the dimension of log_psiw
    jac = tape.jacobian(log_psiw, _wavefunction.trainable_variables)

    return jac
@tf.function
def jacobian_grad(inputs, _wavefunction):
    
    n_walkers = inputs.shape[0]
    
    with tf.GradientTape(persistent=True) as tape:
        log_psiw = _wavefunction(inputs)

        split = tf.split(log_psiw, n_walkers)

    # print(split)
    # By default, this essentially SUMS over the dimension of log_psiw
    grad = [tape.gradient(s, _wavefunction.trainable_variables) for s in split]

    jac = []
    for i, l in enumerate(_wavefunction.trainable_variables):
        temp = tf.stack([g[i] for g in grad])
        temp = tf.reshape(temp,  log_psiw.shape + l.shape)
        jac.append(temp)
    
    return jac
                             
start = time.time()
jc = jacobian_comp(x_input, wavefunction)
print("Compilation time: ", time.time() - start)
start = time.time()
jg = jacobian_grad(x_input, wavefunction)
print("Compilation time: ", time.time() - start)
%timeit jacobian_comp(x_input, wavefunction)
%timeit jacobian_grad(x_input, wavefunction)
print(jg[0].shape) 
print(jc[0].shape)<jupyter_output>(100, 1, 3, 8)
(100, 1, 3, 8)
{'fraction_non_alphanumeric': 0.1061835478455512, 'fraction_numerical': 0.023642977056519306, 'mean_word_length': 3.068867387592487, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '291885', 'license': 'MIT', 'n_tokens_mistral': 2523, 'n_tokens_neox': 2329, 'n_words': 598, 'path': 'Tensorflow Jacobians.ipynb', 'repo_name': 'coreyjadams/AI4NP_School', 'size': '7148'}