# Julia Equivalent to PyTorch's \`torch.nn.functional.affine\_grid\`

**URL:** <https://discourse.julialang.org/t/julia-equivalent-to-pytorchs-torch-nn-functional-affine-grid/118205>\
**Category:** General Usage\
**Created:** [August 14, 2024, 11:18pm UTC](https://discourse.julialang.org/t/julia-equivalent-to-pytorchs-torch-nn-functional-affine-grid/118205 "2024-08-14T23:18:37Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![Dale\_James\_Black](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dale_james_black/32/11595_2.png) [@Dale\_James\_Black](https://discourse.julialang.org/u/Dale_James_Black)\
**Post date:** [August 14, 2024, 11:18pm UTC](https://discourse.julialang.org/t/julia-equivalent-to-pytorchs-torch-nn-functional-affine-grid/118205/1 "2024-08-14T23:18:37Z")

</div>

I asked on Slack without much help. Is there an equivalent to the `affine_grid` function from Pytorch?

This numpy version will work if I want to port it to Julia but would rather see if any Julia packages provide a more robust version of this similar to Pytorch’s. I havent found one yet

```python
def affine_grid_3d(theta, shape, align_corners=False):
    corners = np.ones(3) if align_corners else 1 - 1 / np.array(shape)
    grid_bins = [np.linspace(-c, c, s, dtype=theta.dtype) for c, s in zip(corners, shape)]
    
    x_grid, y_grid, z_grid = np.meshgrid(grid_bins[0], grid_bins[1], grid_bins[2], indexing='ij', copy=False)
    
    grid = np.stack([z_grid, y_grid, x_grid, np.ones_like(x_grid)])
    grid = grid.reshape(4, -1)
    grid = grid[None].repeat(len(theta), axis=0)
    grid = theta @ grid
    grid = grid.reshape(len(theta), 3, *shape)
    return grid.transpose(0, 2, 3, 4, 1)

    return grid

```