[tensor] add get_subtensor

diff --git a/src/ai/tensor.c b/src/ai/tensor.c
index 8a7b76d..efa4210 100644 (file)
--- a/src/ai/tensor.c
+++ b/src/ai/tensor.c
@@ -80,6 +80,36 @@ uint_t aubio_fmat_as_tensor(fmat_t *o, aubio_tensor_t *c) {
   return AUBIO_OK;
 }
 
+uint_t aubio_tensor_get_subtensor(aubio_tensor_t *t, uint_t i,
+    aubio_tensor_t *st)
+{
+  uint_t j;
+  if (!t || !st) return AUBIO_FAIL;
+  if (i >= t->shape[0]) {
+    AUBIO_ERR("tensor: index %d out of range, only %d subtensors\n",
+        i, t->shape[0]);
+    return AUBIO_FAIL;
+  }
+  if(t->ndim > 1) {
+    st->ndim = t->ndim - 1;
+    for (j = 0; j < st->ndim; j++) {
+      st->shape[j] = t->shape[j + 1];
+    }
+    for (j = st->ndim; j < AUBIO_TENSOR_MAXDIM; j++) {
+      st->shape[j] = 0;
+    }
+    st->size = t->size / t->shape[0];
+  } else {
+    st->ndim = 1;
+    st->shape[0] = 1;
+    st->size = 1;
+  }
+  // st was allocated on the stack, row indices are lost
+  st->data = NULL;
+  st->buffer = &t->buffer[0] + st->size * i;
+  return AUBIO_OK;
+}
+
 smpl_t aubio_tensor_max(aubio_tensor_t *t)
 {
   uint_t i;

diff --git a/src/ai/tensor.h b/src/ai/tensor.h
index 656b7fa..181a853 100644 (file)
--- a/src/ai/tensor.h
+++ b/src/ai/tensor.h
@@ -58,6 +58,9 @@ uint_t aubio_fvec_as_tensor(fvec_t *o, aubio_tensor_t *c);
 uint_t aubio_tensor_as_fmat(aubio_tensor_t *c, fmat_t *o);
 uint_t aubio_fmat_as_tensor(fmat_t *o, aubio_tensor_t *c);
 
+uint_t aubio_tensor_get_subtensor(aubio_tensor_t *t, uint_t i,
+        aubio_tensor_t *st);
+
 smpl_t aubio_tensor_max(aubio_tensor_t *t);
 
 #define AUBIO_ASSERT_EQUAL_SHAPE(t1, t2) { \