Another incarnation of `expand_dims`, following the contract https://pytorch.org/docs/stable/generated/torch.unsqueeze.html (in particular, the `axis` argument can be negative).
(x: Node, axis: i64)
| 276 | // Another incarnation of `expand_dims`, following the contract https://pytorch.org/docs/stable/generated/torch.unsqueeze.html |
| 277 | // (in particular, the `axis` argument can be negative). |
| 278 | pub fn unsqueeze(x: Node, axis: i64) -> Result<Node> { |
| 279 | let (mut shape, sc) = match x.get_type()? { |
| 280 | Type::Array(shape, sc) => (shape, sc), |
| 281 | Type::Scalar(sc) => (vec![], sc), |
| 282 | _ => { |
| 283 | return Err(runtime_error!( |
| 284 | "Expected array or scalar type, got {:?}", |
| 285 | x.get_type()? |
| 286 | )) |
| 287 | } |
| 288 | }; |
| 289 | if axis < -(shape.len() as i64) - 1 || axis > shape.len() as i64 { |
| 290 | return Err(runtime_error!( |
| 291 | "Expected axis in range [{}, {}], got {}", |
| 292 | -(shape.len() as i64) - 1, |
| 293 | shape.len() as i64, |
| 294 | axis |
| 295 | )); |
| 296 | } |
| 297 | let pos = if axis < 0 { |
| 298 | shape.len() as i64 + axis + 1 |
| 299 | } else { |
| 300 | axis |
| 301 | }; |
| 302 | shape.insert(pos as usize, 1); |
| 303 | x.reshape(array_type(shape, sc)) |
| 304 | } |
| 305 | |
| 306 | /// Similar to `Sum`, but for multiplication. Reduces over the last dimension. |
| 307 | /// The use-case where it makes the most sense is when the type is BIT. |
no test coverage detected