I want to create a tensor with `torch.zeros`

based on the shape of an input to the function. Then I want to vectorize the function with `torch.vmap`

.

Something like this:

```
poly_batched = torch.tensor([[1, 2, 3, 4], [1, 2, 3, 4]])
def polycompanion(polynomial):
deg = polynomial.shape[-1] - 2
companion = torch.zeros((deg+1, deg+1))
companion[1:,:-1] = torch.eye(deg)
companion[:,-1] = -1. * polynomial[:-1] / polynomial[-1]
return companion
polycompanion_vmap = torch.vmap(polycompanion)
print(polycompanion_vmap(poly_batched))
```

The problem is that the batched version will not work, because `companion`

won’t be a `BatchedTensor`

, unlike `polynomial`

, which was the input.

There is a workaround:

```
poly_batched = torch.tensor([[1, 2, 3, 4], [1, 2, 3, 4]])
def polycompanion(polynomial,companion):
deg = companion.shape[-1] - 1
companion[1:,:-1] = torch.eye(deg)
companion[:,-1] = -1. * polynomial[:-1] / polynomial[-1]
return companion
polycompanion_vmap = torch.vmap(polycompanion)
print(polycompanion_vmap(poly_batched, torch.zeros(poly_batched.shape[0],poly_batched.shape[-1]-1, poly_batched.shape[-1]-1)))
```

Output:

```
tensor([[[ 0.0000, 0.0000, -0.2500],
[ 1.0000, 0.0000, -0.5000],
[ 0.0000, 1.0000, -0.7500]],
[[ 0.0000, 0.0000, -0.2500],
[ 1.0000, 0.0000, -0.5000],
[ 0.0000, 1.0000, -0.7500]]])
```

But this is ugly.

Is there a solution for this? Will this be supported in the future?

Note: If you use `torch.zeros_like`

on an input to the function it works and creates `BatchedTensor`

but this doesn’t help me here.

Thanks in advance for the help!