


assert feature_vec.dim() == 1
return feature_vec.dot(weights).relu()

def compute_loss(weights, example, target):
    y = model(weights, example)
    return ((y - target) ** 2).mean()  # MSELoss
weights = torch.randn(feature_size, requires_grad=True)
examples = torch.randn(batch_size, feature_size)
targets = torch.randn(batch_size)
inputs = (weights, examples, targets)
grad_weight_per_example = vmap(grad(compute_loss), in_dims=(None, 0, 0))(*inputs)



Example of using "grad" with "has_aux" and "argnums":



from torch.func import grad
def my_loss_func(y, y_pred):
   loss_per_sample = (0.5 * y_pred - y) ** 2
   loss = loss_per_sample.mean()
   return loss, (y_pred, loss_per_sample)
fn = grad(my_loss_func, argnums=(0, 1), has_aux=True)
y_true = torch.rand(4)
y_preds = torch.rand(4, requires_grad=True)
out = fn(y_true, y_preds)


