// Supports arbitrary batch dimensions for self and A
std::tuple<Tensor,Tensor> gesv(const Tensor& self, const Tensor& A) {
if (self.dim() <= 2 && A.dim() <= 2) {
// TODO: #7102: It's not necessary to have gesv (single) bindings for both
// TH and ATen. We should remove the TH gesv bindings, especially
// since the lapackGesv function is already in ATen.
return at::_gesv_single(self, A);
}
checkInputs(self, A);
// broadcast the batch dimensions of self and A.
IntList self_batch_sizes(self.sizes().data(), self.ndimension() - 2);
IntList A_batch_sizes(A.sizes().data(), A.ndimension() - 2);
std::vector<int64_t> expand_batch_portion =
infer_size(self_batch_sizes, A_batch_sizes);
std::vector<int64_t> self_expand_size({expand_batch_portion});
self_expand_size.insert(self_expand_size.end(),
{ self.size(-2), self.size(-1) });
std::vector<int64_t> A_expand_size({expand_batch_portion});
A_expand_size.insert(A_expand_size.end(),
{ A.size(-2), A.size(-1) });
Tensor self_broadcasted = self.expand(self_expand_size);
Tensor A_broadcasted = A.expand(A_expand_size);
return self.type()._gesv_helper(self_broadcasted, A_broadcasted);
}
static int img_read_packet(AVFormatContext *s1, AVPacket *pkt)
{
VideoData *s = s1->priv_data;
char filename[1024];
int i;
int size[3]={0}, ret[3]={0};
ByteIOContext f1[3], *f[3]= {&f1[0], &f1[1], &f1[2]};
AVCodecContext *codec= s1->streams[0]->codec;
if (!s->is_pipe) {
/* loop over input */
if (s1->loop_input && s->img_number > s->img_last) {
s->img_number = s->img_first;
}
if (av_get_frame_filename(filename, sizeof(filename),
s->path, s->img_number)<0 && s->img_number > 1)
return AVERROR_IO;
for(i=0; i<3; i++){
if (url_fopen(f[i], filename, URL_RDONLY) < 0)
return AVERROR_IO;
size[i]= url_fsize(f[i]);
if(codec->codec_id != CODEC_ID_RAWVIDEO)
break;
filename[ strlen(filename) - 1 ]= 'U' + i;
}
if(codec->codec_id == CODEC_ID_RAWVIDEO && !codec->width)
infer_size(&codec->width, &codec->height, size[0]);
} else {
f[0] = &s1->pb;
if (url_feof(f[0]))
return AVERROR_IO;
size[0]= 4096;
}
av_new_packet(pkt, size[0] + size[1] + size[2]);
pkt->stream_index = 0;
pkt->flags |= PKT_FLAG_KEY;
pkt->size= 0;
for(i=0; i<3; i++){
if(size[i]){
ret[i]= get_buffer(f[i], pkt->data + pkt->size, size[i]);
if (!s->is_pipe)
url_fclose(f[i]);
if(ret[i]>0)
pkt->size += ret[i];
}
}
if (ret[0] <= 0 || ret[1]<0 || ret[2]<0) {
av_free_packet(pkt);
return AVERROR_IO; /* signal EOF */
} else {
s->img_count++;
s->img_number++;
return 0;
}
}
static int yuv_read(ByteIOContext *f,
int (*alloc_cb)(void *opaque, AVImageInfo *info), void *opaque)
{
ByteIOContext pb1, *pb = &pb1;
int img_size, ret;
char fname[1024], *p;
int size;
URLContext *h;
AVImageInfo info1, *info = &info1;
img_size = url_fsize(f);
/* XXX: hack hack */
h = url_fileno(f);
url_get_filename(h, fname, sizeof(fname));
if (infer_size(&info->width, &info->height, img_size) < 0) {
return AVERROR_IO;
}
info->pix_fmt = PIX_FMT_YUV420P;
ret = alloc_cb(opaque, info);
if (ret)
return ret;
size = info->width * info->height;
p = strrchr(fname, '.');
if (!p || p[1] != 'Y')
return AVERROR_IO;
get_buffer(f, info->pict.data[0], size);
p[1] = 'U';
if (url_fopen(pb, fname, URL_RDONLY) < 0)
return AVERROR_IO;
get_buffer(pb, info->pict.data[1], size / 4);
url_fclose(pb);
p[1] = 'V';
if (url_fopen(pb, fname, URL_RDONLY) < 0)
return AVERROR_IO;
get_buffer(pb, info->pict.data[2], size / 4);
url_fclose(pb);
return 0;
}
int ff_img_read_packet(AVFormatContext *s1, AVPacket *pkt)
{
VideoDemuxData *s = s1->priv_data;
char filename_bytes[1024];
char *filename = filename_bytes;
int i, res;
int size[3] = { 0 }, ret[3] = { 0 };
AVIOContext *f[3] = { NULL };
AVCodecParameters *par = s1->streams[0]->codecpar;
if (!s->is_pipe) {
/* loop over input */
if (s->loop && s->img_number > s->img_last) {
s->img_number = s->img_first;
}
if (s->img_number > s->img_last)
return AVERROR_EOF;
if (s->pattern_type == PT_NONE) {
av_strlcpy(filename_bytes, s->path, sizeof(filename_bytes));
} else if (s->use_glob) {
#if HAVE_GLOB
filename = s->globstate.gl_pathv[s->img_number];
#endif
} else {
if (av_get_frame_filename(filename_bytes, sizeof(filename_bytes),
s->path,
s->img_number) < 0 && s->img_number > 1)
return AVERROR(EIO);
}
for (i = 0; i < 3; i++) {
if (s1->pb &&
!strcmp(filename_bytes, s->path) &&
!s->loop &&
!s->split_planes) {
f[i] = s1->pb;
} else if (s1->io_open(s1, &f[i], filename, AVIO_FLAG_READ, NULL) < 0) {
if (i >= 1)
break;
av_log(s1, AV_LOG_ERROR, "Could not open file : %s\n",
filename);
return AVERROR(EIO);
}
size[i] = avio_size(f[i]);
if (!s->split_planes)
break;
filename[strlen(filename) - 1] = 'U' + i;
}
if (par->codec_id == AV_CODEC_ID_NONE) {
AVProbeData pd = { 0 };
AVInputFormat *ifmt;
uint8_t header[PROBE_BUF_MIN + AVPROBE_PADDING_SIZE];
int ret;
int score = 0;
ret = avio_read(f[0], header, PROBE_BUF_MIN);
if (ret < 0)
return ret;
memset(header + ret, 0, sizeof(header) - ret);
avio_skip(f[0], -ret);
pd.buf = header;
pd.buf_size = ret;
pd.filename = filename;
ifmt = av_probe_input_format3(&pd, 1, &score);
if (ifmt && ifmt->read_packet == ff_img_read_packet && ifmt->raw_codec_id)
par->codec_id = ifmt->raw_codec_id;
}
if (par->codec_id == AV_CODEC_ID_RAWVIDEO && !par->width)
infer_size(&par->width, &par->height, size[0]);
} else {
f[0] = s1->pb;
if (avio_feof(f[0]) && s->loop && s->is_pipe)
avio_seek(f[0], 0, SEEK_SET);
if (avio_feof(f[0]))
return AVERROR_EOF;
if (s->frame_size > 0) {
size[0] = s->frame_size;
} else if (!s1->streams[0]->parser) {
size[0] = avio_size(s1->pb);
} else {
size[0] = 4096;
}
}
res = av_new_packet(pkt, size[0] + size[1] + size[2]);
if (res < 0) {
goto fail;
}
pkt->stream_index = 0;
pkt->flags |= AV_PKT_FLAG_KEY;
if (s->ts_from_file) {
struct stat img_stat;
if (stat(filename, &img_stat)) {
res = AVERROR(EIO);
goto fail;
}
pkt->pts = (int64_t)img_stat.st_mtime;
#if HAVE_STRUCT_STAT_ST_MTIM_TV_NSEC
if (s->ts_from_file == 2)
pkt->pts = 1000000000*pkt->pts + img_stat.st_mtim.tv_nsec;
#endif
av_add_index_entry(s1->streams[0], s->img_number, pkt->pts, 0, 0, AVINDEX_KEYFRAME);
} else if (!s->is_pipe) {
pkt->pts = s->pts;
}
if (s->is_pipe)
pkt->pos = avio_tell(f[0]);
pkt->size = 0;
for (i = 0; i < 3; i++) {
if (f[i]) {
ret[i] = avio_read(f[i], pkt->data + pkt->size, size[i]);
if (s->loop && s->is_pipe && ret[i] == AVERROR_EOF) {
if (avio_seek(f[i], 0, SEEK_SET) >= 0) {
pkt->pos = 0;
ret[i] = avio_read(f[i], pkt->data + pkt->size, size[i]);
}
}
if (!s->is_pipe && f[i] != s1->pb)
ff_format_io_close(s1, &f[i]);
if (ret[i] > 0)
pkt->size += ret[i];
}
}
if (ret[0] <= 0 || ret[1] < 0 || ret[2] < 0) {
av_packet_unref(pkt);
if (ret[0] < 0) {
res = ret[0];
} else if (ret[1] < 0) {
res = ret[1];
} else if (ret[2] < 0) {
res = ret[2];
} else {
res = AVERROR_EOF;
}
goto fail;
} else {
s->img_count++;
s->img_number++;
s->pts++;
return 0;
}
fail:
if (!s->is_pipe) {
for (i = 0; i < 3; i++) {
if (f[i] != s1->pb)
ff_format_io_close(s1, &f[i]);
}
}
return res;
}
static int img_read_packet(AVFormatContext *s1, AVPacket *pkt)
{
VideoDemuxData *s = s1->priv_data;
char filename[1024];
int i, res;
int size[3] = { 0 }, ret[3] = { 0 };
AVIOContext *f[3] = { NULL };
AVCodecContext *codec = s1->streams[0]->codec;
if (!s->is_pipe) {
/* loop over input */
if (s->loop && s->img_number > s->img_last) {
s->img_number = s->img_first;
}
if (s->img_number > s->img_last)
return AVERROR_EOF;
if (av_get_frame_filename(filename, sizeof(filename),
s->path,
s->img_number) < 0 && s->img_number > 1)
return AVERROR(EIO);
for (i = 0; i < 3; i++) {
if (s1->io_open(s1, &f[i], filename, AVIO_FLAG_READ, NULL) < 0) {
if (i >= 1)
break;
av_log(s1, AV_LOG_ERROR, "Could not open file : %s\n",
filename);
return AVERROR(EIO);
}
size[i] = avio_size(f[i]);
if (codec->codec_id != AV_CODEC_ID_RAWVIDEO)
break;
filename[strlen(filename) - 1] = 'U' + i;
}
if (codec->codec_id == AV_CODEC_ID_RAWVIDEO && !codec->width)
infer_size(&codec->width, &codec->height, size[0]);
} else {
f[0] = s1->pb;
if (f[0]->eof_reached)
return AVERROR(EIO);
size[0] = 4096;
}
res = av_new_packet(pkt, size[0] + size[1] + size[2]);
if (res < 0)
return res;
pkt->stream_index = 0;
pkt->flags |= AV_PKT_FLAG_KEY;
pkt->size = 0;
for (i = 0; i < 3; i++) {
if (f[i]) {
ret[i] = avio_read(f[i], pkt->data + pkt->size, size[i]);
if (!s->is_pipe)
ff_format_io_close(s1, &f[i]);
if (ret[i] > 0)
pkt->size += ret[i];
}
}
if (ret[0] <= 0 || ret[1] < 0 || ret[2] < 0) {
av_packet_unref(pkt);
return AVERROR(EIO); /* signal EOF */
} else {
s->img_count++;
s->img_number++;
return 0;
}
}
/*
Matrix product of two Tensors.
The behavior depends on the dimensionality of the Tensors as follows:
- If both Tensors are 1-dimensional, the dot product (scalar) is returned.
- If both arguments are 2-dimensional, the matrix-matrix product is returned.
- If the first argument is 1-dimensional and the second argument is 2-dimensional,
a 1 is prepended to its dimension for the purpose of the matrix multiply.
After the matrix multiply, the prepended dimension is removed.
- If the first argument is 2-dimensional and the second argument is 1-dimensional,
the matrix-vector product is returned.
- If both arguments are at least 1-dimensional and at least one argument is
N-dimensional (where N > 2), then a batched matrix multiply is returned. If the first
argument is 1-dimensional, a 1 is prepended to its dimension for the purpose of the
batched matrix multiply and removed after. If the second argument is 1-dimensional, a
1 is appended to its dimension for the purpose of the batched matrix multiple and removed after.
The non-matrix (i.e. batch) dimensions are broadcasted (and thus
must be broadcastable). For example, if tensor1 is a (j x 1 x n x m) Tensor
and tensor2 is a (k x m x p) Tensor, the returned tensor will be an (j x k x n x p) Tensor.
*/
Tensor matmul(const Tensor & tensor1, const Tensor & tensor2) {
auto dim_tensor1 = tensor1.dim();
auto dim_tensor2 = tensor2.dim();
if (dim_tensor1 == 1 && dim_tensor2 == 1) {
return tensor1.dot(tensor2);
} else if (dim_tensor1 == 2 && dim_tensor2 == 1) {
return tensor1.mv(tensor2);
} else if (dim_tensor1 == 1 && dim_tensor2 == 2) {
return tensor1.unsqueeze(0).mm(tensor2).squeeze_(0);
} else if (dim_tensor1 == 2 && dim_tensor2 == 2) {
return tensor1.mm(tensor2);
} else if (dim_tensor1 >= 3 && (dim_tensor2 == 1 || dim_tensor2 == 2)) {
// optimization: use mm instead of bmm by folding tensor1's batch into
// its leading matrix dimension.
Tensor t2 = dim_tensor2 == 1 ? tensor2.unsqueeze(-1) : tensor2;
auto size1 = tensor1.sizes();
auto size2 = t2.sizes();
std::vector<int64_t> output_size;
output_size.insert(output_size.end(), size1.begin(), size1.end() - 1);
if (dim_tensor2 > 1) {
output_size.push_back(size2[dim_tensor2 - 1]);
}
// fold the batch into the first dimension
Tensor t1 = tensor1.contiguous().view({-1, size1[size1.size() - 1]});
return at::_unsafe_view(t1.mm(t2), output_size);
} else if ((dim_tensor1 >= 1 && dim_tensor2 >= 1) && (dim_tensor1 >= 3 || dim_tensor2 >= 3)) {
// We are multiplying b1 x n x m1 by x2 x m2 x p (where b1 can be a list);
// we track m1 vs m2 separately even though they must match for nicer error messages
int64_t n = dim_tensor1 > 1 ? tensor1.size(-2) : 1;
int64_t m1 = tensor1.size(-1);
IntList batch_tensor1(tensor1.sizes().data(), std::max<int64_t>(dim_tensor1 - 2, 0));
int64_t m2 = dim_tensor2 > 1 ? tensor2.size(-2) : 1;
int64_t p = tensor2.size(-1);
IntList batch_tensor2(tensor2.sizes().data(), std::max<int64_t>(dim_tensor2 - 2, 0));
// expand the batch portion (i.e. cut off matrix dimensions and expand rest)
std::vector<int64_t> expand_batch_portion = infer_size(batch_tensor1, batch_tensor2);
std::vector<int64_t> tensor1_expand_size(expand_batch_portion);
tensor1_expand_size.insert(tensor1_expand_size.end(), {n, m1});
std::vector<int64_t> tensor2_expand_size(expand_batch_portion);
tensor2_expand_size.insert(tensor2_expand_size.end(), {m2, p});
int expand_batch_product = std::accumulate(expand_batch_portion.begin(), expand_batch_portion.end(),
1, std::multiplies<int64_t>());
std::vector<int64_t> tensor1_bmm_view({expand_batch_product});
tensor1_bmm_view.insert(tensor1_bmm_view.end(), {n, m1});
std::vector<int64_t> tensor2_bmm_view({expand_batch_product});
tensor2_bmm_view.insert(tensor2_bmm_view.end(), {m2, p});
// flatten expanded batches
Tensor tensor1_expanded = tensor1.expand(tensor1_expand_size).contiguous().view(tensor1_bmm_view);
Tensor tensor2_expanded = tensor2.expand(tensor2_expand_size).contiguous().view(tensor2_bmm_view);
Tensor output = tensor1_expanded.bmm(tensor2_expanded);
// reshape batches back into result
std::vector<int64_t> output_shape(expand_batch_portion);
if (dim_tensor1 > 1) {
output_shape.push_back(n);
}
if (dim_tensor2 > 1) {
output_shape.push_back(p);
}
return at::_unsafe_view(output, output_shape);
}
runtime_error("both arguments to matmul need to be at least 1D, but they are %dD and %dD",
dim_tensor1, dim_tensor2);
}