# Speeding up per-sample gradients?

**URL:** https://discourse.julialang.org/t/speeding-up-per-sample-gradients/109322
**Category:** Machine Learning
**Tags:** question, autodiff
**Created:** [January 26, 2024, 10:45pm UTC](https://discourse.julialang.org/t/speeding-up-per-sample-gradients/109322 "2024-01-26T22:45:58Z")
**Posts on this page:** 1
**Showing post:** 9

<div class="post-metadata">

### Author: ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [January 30, 2024, 4:31pm UTC](https://discourse.julialang.org/t/speeding-up-per-sample-gradients/109322/9 "2024-01-30T16:31:27Z")

</div>

Mostly. Just broadcasting won’t work as well as `vmap`, however, because some of the operations being broadcasted are already vectorized (e.g. BLAS). `vmap` will actually modify those calls (using dispatch in PyTorch and source code transforms in JAX) to use batched implementations whenever it encounters them.

Now that said, some functions and Flux layers are flexible enough to already work for this without a `vmap`-like treatment. See this topic posted about a month ago: [Flux loss with contribution gradient is slow - #5 by Jonas208](https://discourse.julialang.org/t/flux-loss-with-contribution-gradient-is-slow/107956/5). Basically, changing your loss function to compute a loss for each sample individually and then summing should be enough for a MLP. In fact, `mean_batch_grad` and `map_grad` currently return the exact same gradients because `sum(map(x -> sum(model(xs)), xs)) == sum(model(xs))`!

---

_[View the full topic](https://discourse.julialang.org/t/speeding-up-per-sample-gradients/109322)._