forked from M-Labs/nac3
core/ndstrides: add NDArrayObject::atleast_nd
This commit is contained in:
parent
b175409a9b
commit
ea1410f9ff
|
@ -2,6 +2,7 @@ pub mod factory;
|
|||
pub mod indexing;
|
||||
pub mod nditer;
|
||||
pub mod shape_util;
|
||||
pub mod view;
|
||||
|
||||
use inkwell::{
|
||||
context::Context,
|
||||
|
|
|
@ -0,0 +1,29 @@
|
|||
use crate::codegen::{CodeGenContext, CodeGenerator};
|
||||
|
||||
use super::{indexing::RustNDIndex, NDArrayObject};
|
||||
|
||||
impl<'ctx> NDArrayObject<'ctx> {
|
||||
/// Make sure the ndarray is at least `ndmin`-dimensional.
|
||||
///
|
||||
/// If this ndarray's `ndims` is less than `ndmin`, a view is created on this with 1s prepended to the shape.
|
||||
/// If this ndarray's `ndims` is not less than `ndmin`, this function does nothing and return this ndarray.
|
||||
#[must_use]
|
||||
pub fn atleast_nd<G: CodeGenerator + ?Sized>(
|
||||
&self,
|
||||
generator: &mut G,
|
||||
ctx: &mut CodeGenContext<'ctx, '_>,
|
||||
ndmin: u64,
|
||||
) -> Self {
|
||||
if self.ndims < ndmin {
|
||||
// return this_ndarray[np.newaxis, np.newaxis, and more, ...]
|
||||
let mut indices = vec![];
|
||||
for _ in self.ndims..ndmin {
|
||||
indices.push(RustNDIndex::NewAxis);
|
||||
}
|
||||
indices.push(RustNDIndex::Ellipsis);
|
||||
self.index(generator, ctx, &indices)
|
||||
} else {
|
||||
*self
|
||||
}
|
||||
}
|
||||
}
|
Loading…
Reference in New Issue