| 935 | } |
| 936 | |
| 937 | Tensor Tensor::SubSlice(int64 index) const { |
| 938 | CHECK_GE(dims(), 1); // Crash ok. |
| 939 | CHECK_LE(0, index); // Crash ok. |
| 940 | int64 dim0_size = shape_.dim_size(0); |
| 941 | CHECK_LE(index, dim0_size); // Crash ok. |
| 942 | Tensor ret; |
| 943 | ret.shape_ = shape_; |
| 944 | ret.shape_.RemoveDim(0); |
| 945 | ret.set_dtype(dtype()); |
| 946 | ret.buf_ = nullptr; |
| 947 | if (dim0_size > 0) { |
| 948 | const int64 elems_per_dim0 = NumElements() / dim0_size; |
| 949 | const int64 delta = index * elems_per_dim0; |
| 950 | const int64 num_elems = elems_per_dim0; |
| 951 | if (buf_) { |
| 952 | DataType dt = dtype(); |
| 953 | CASES(dt, ret.buf_ = new SubBuffer<T>(buf_, delta, num_elems)); |
| 954 | } |
| 955 | } |
| 956 | return ret; |
| 957 | } |
| 958 | |
| 959 | bool Tensor::FromProto(const TensorProto& proto) { |
| 960 | return FromProto(get_default_cpu_allocator(), proto); |