Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[LLVM Integrate] Bump {stablehlo, torch-mlir, llvm} #73

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions deps.bzl
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,8 @@ def third_party_deps():
path = local_llvm_repo_path(),
)
else:
LLVM_COMMIT = "a952c123880eb1168f1021b116485e27170d48ca"
LLVM_SHA256 = "c33580d7581abece69ea20d640a9e49b96472633611025663313cf0644567fa2"
LLVM_COMMIT = "5207632f8698a2fab0c4cdcdf2f7ad9aaf96e06f"
LLVM_SHA256 = "02b8d285091e2f7995dd4f339d5312f935966c085e4b32458ef2b85c43e1d237"
http_archive(
name = "llvm-raw",
build_file_content = "# empty",
Expand All @@ -39,8 +39,8 @@ def third_party_deps():
path = local_torch_mlir_repo_path(),
)
else:
TORCH_MLIR_COMMIT = "45eaeaaf3640f1f7f7b4743cf5844c3b9d519fdf"
TORCH_MLIR_SHA256 = "c68ef55383e02855d2dbe769470ca27966387505f6b7c9aff667068757636786"
TORCH_MLIR_COMMIT = "c7d52f63b482b2c30f4efb435ce0cc2efeab25d9"
TORCH_MLIR_SHA256 = "db037fac59f750340798a8680b19c79c3e950ddbecc52436f53dca340cb0352a"
http_archive(
name = "torch-mlir-raw",
build_file_content = "# empty",
Expand All @@ -55,8 +55,8 @@ def third_party_deps():
path = local_stablehlo_repo_path(),
)
else:
STABLEHLO_COMMIT = "271e8634de184fbfafd677d3876170feb6d08c97"
STABLEHLO_SHA256 = "0d567e9a5f0f6c8487272dc27539403f1c9cae5039e64c52a648ba08ae11f86d"
STABLEHLO_COMMIT = "d41390c3a731ba038e6363f75fcd135e6f727039"
STABLEHLO_SHA256 = "f278ef6aa908a9b8be482481dd408a7d6887a5117a4d425a0cb4a95a242585cc"
http_archive(
name = "stablehlo",
sha256 = STABLEHLO_SHA256,
Expand Down
29 changes: 20 additions & 9 deletions lib/Pipeline/Pipeline.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,15 @@

#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
#include "mlir/Conversion/BufferizationToMemRef/BufferizationToMemRef.h"
#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
#include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
#include "mlir/Conversion/MathToLibm/MathToLibm.h"
#include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
#include "mlir/Dialect/Arith/Transforms/Passes.h"
#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Func/Transforms/Passes.h"
Expand Down Expand Up @@ -71,17 +73,26 @@ static void createTcpToLlvmPipeline(OpPassManager &pm) {
pm.addNestedPass<func::FuncOp>(tcp::createConvertTcpToTensorPass());
pm.addNestedPass<func::FuncOp>(tcp::createConvertTcpToArithPass());

// Bufferize tensor -> memref.
pm.addNestedPass<func::FuncOp>(
bufferization::createEmptyTensorToAllocTensorPass());
pm.addNestedPass<func::FuncOp>(
bufferization::createBufferizationBufferizePass());
pm.addNestedPass<func::FuncOp>(createLinalgBufferizePass());
pm.addNestedPass<func::FuncOp>(tensor::createTensorBufferizePass());
pm.addPass(arith::createConstantBufferizePass());
pm.addPass(func::createFuncBufferizePass());
// One-shot bufferize tensor -> memref, from
// https://mlir.llvm.org/docs/Bufferization/.
bufferization::OneShotBufferizationOptions bufferizationOptions;
bufferizationOptions.bufferizeFunctionBoundaries = true;
bufferizationOptions.setFunctionBoundaryTypeConversion(
bufferization::LayoutMapOption::IdentityLayoutMap);
pm.addPass(bufferization::createOneShotBufferizePass(bufferizationOptions));
pm.addNestedPass<func::FuncOp>(
bufferization::createFinalizingBufferizePass());
// Buffer deallocation pipeline for automatically inserting
// buffer deallocation ops after one-shot bufferization.
// https://sourcegraph.com/github.com/llvm/llvm-project@09bc1e825068f314db71ee7eb32d9f93c5ac87a0/-/blob/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp?L21
pm.addPass(createCanonicalizerPass());
pm.addPass(bufferization::createOwnershipBasedBufferDeallocationPass());
pm.addPass(createCanonicalizerPass());
pm.addPass(bufferization::createBufferDeallocationSimplificationPass());
pm.addPass(bufferization::createLowerDeallocationsPass());
pm.addPass(createCSEPass());
pm.addPass(createCanonicalizerPass());
pm.addPass(createBufferizationToMemRefPass());

// Blanket-convert any remaining linalg ops to loops if any remain.
pm.addNestedPass<func::FuncOp>(createConvertLinalgToLoopsPass());
Expand Down
131 changes: 70 additions & 61 deletions requirements_lock.txt
Original file line number Diff line number Diff line change
Expand Up @@ -6,19 +6,19 @@
#
--find-links https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels

filelock==3.14.0 \
--hash=sha256:43339835842f110ca7ae60f1e1c160714c5a6afd15a2873419ab185334975c0f \
--hash=sha256:6ea72da3be9b8c82afd3edcf99f2fffbb5076335a5ae4d03248bb5b6c3eae78a
filelock==3.15.3 \
--hash=sha256:0151273e5b5d6cf753a61ec83b3a9b7d8821c39ae9af9d7ecf2f9e2f17404103 \
--hash=sha256:e1199bf5194a2277273dacd50269f0d87d0682088a3c561c15674ea9005d8635
# via
# torch
# triton
fsspec==2024.3.1 \
--hash=sha256:918d18d41bf73f0e2b261824baeb1b124bcf771767e3a26425cd7dec3332f512 \
--hash=sha256:f39780e282d7d117ffb42bb96992f8a90795e4d0fb0f661a70ca39fe9c43ded9
fsspec==2024.6.0 \
--hash=sha256:58d7122eb8a1a46f7f13453187bfea4972d66bf01618d37366521b1998034cee \
--hash=sha256:f579960a56e6d8038a9efc8f9c77279ec12e6299aa86b0769a7e9c46b94527c2
# via torch
jinja2==3.1.3 \
--hash=sha256:7d6d50dd97d52cbc355597bd845fabfbac3f551e1f99619e39a35ce8c370b5fa \
--hash=sha256:ac8bd6544d4bb2c9792bf3a159e80bba8fda7f07e81bc3aed565432d5925ba90
jinja2==3.1.4 \
--hash=sha256:4a3aee7acbbe7303aede8e9648d13b8bf88a429282aa6122a993f0ac800cb369 \
--hash=sha256:bc5dd2abb727a5319567b7a813e6a2e7318c39f4f487cfe6c89c6f9c7d25197d
# via torch
markupsafe==2.1.5 \
--hash=sha256:00e046b6dd71aa03a41079792f8473dc494d564611a8f89bbbd7cb93295ebdcf \
Expand Down Expand Up @@ -90,65 +90,74 @@ networkx==3.3 \
--hash=sha256:0c127d8b2f4865f59ae9cb8aafcd60b5c70f3241ebd66f7defad7c4ab90126c9 \
--hash=sha256:28575580c6ebdaf4505b22c6256a2b9de86b316dc63ba9e93abde3d78dfdbcf2
# via torch
numpy==1.26.4 \
--hash=sha256:03a8c78d01d9781b28a6989f6fa1bb2c4f2d51201cf99d3dd875df6fbd96b23b \
--hash=sha256:08beddf13648eb95f8d867350f6a018a4be2e5ad54c8d8caed89ebca558b2818 \
--hash=sha256:1af303d6b2210eb850fcf03064d364652b7120803a0b872f5211f5234b399f20 \
--hash=sha256:1dda2e7b4ec9dd512f84935c5f126c8bd8b9f2fc001e9f54af255e8c5f16b0e0 \
--hash=sha256:2a02aba9ed12e4ac4eb3ea9421c420301a0c6460d9830d74a9df87efa4912010 \
--hash=sha256:2e4ee3380d6de9c9ec04745830fd9e2eccb3e6cf790d39d7b98ffd19b0dd754a \
--hash=sha256:3373d5d70a5fe74a2c1bb6d2cfd9609ecf686d47a2d7b1d37a8f3b6bf6003aea \
--hash=sha256:47711010ad8555514b434df65f7d7b076bb8261df1ca9bb78f53d3b2db02e95c \
--hash=sha256:4c66707fabe114439db9068ee468c26bbdf909cac0fb58686a42a24de1760c71 \
--hash=sha256:50193e430acfc1346175fcbdaa28ffec49947a06918b7b92130744e81e640110 \
--hash=sha256:52b8b60467cd7dd1e9ed082188b4e6bb35aa5cdd01777621a1658910745b90be \
--hash=sha256:60dedbb91afcbfdc9bc0b1f3f402804070deed7392c23eb7a7f07fa857868e8a \
--hash=sha256:62b8e4b1e28009ef2846b4c7852046736bab361f7aeadeb6a5b89ebec3c7055a \
--hash=sha256:666dbfb6ec68962c033a450943ded891bed2d54e6755e35e5835d63f4f6931d5 \
--hash=sha256:675d61ffbfa78604709862923189bad94014bef562cc35cf61d3a07bba02a7ed \
--hash=sha256:679b0076f67ecc0138fd2ede3a8fd196dddc2ad3254069bcb9faf9a79b1cebcd \
--hash=sha256:7349ab0fa0c429c82442a27a9673fc802ffdb7c7775fad780226cb234965e53c \
--hash=sha256:7ab55401287bfec946ced39700c053796e7cc0e3acbef09993a9ad2adba6ca6e \
--hash=sha256:7e50d0a0cc3189f9cb0aeb3a6a6af18c16f59f004b866cd2be1c14b36134a4a0 \
--hash=sha256:95a7476c59002f2f6c590b9b7b998306fba6a5aa646b1e22ddfeaf8f78c3a29c \
--hash=sha256:96ff0b2ad353d8f990b63294c8986f1ec3cb19d749234014f4e7eb0112ceba5a \
--hash=sha256:9fad7dcb1aac3c7f0584a5a8133e3a43eeb2fe127f47e3632d43d677c66c102b \
--hash=sha256:9ff0f4f29c51e2803569d7a51c2304de5554655a60c5d776e35b4a41413830d0 \
--hash=sha256:a354325ee03388678242a4d7ebcd08b5c727033fcff3b2f536aea978e15ee9e6 \
--hash=sha256:a4abb4f9001ad2858e7ac189089c42178fcce737e4169dc61321660f1a96c7d2 \
--hash=sha256:ab47dbe5cc8210f55aa58e4805fe224dac469cde56b9f731a4c098b91917159a \
--hash=sha256:afedb719a9dcfc7eaf2287b839d8198e06dcd4cb5d276a3df279231138e83d30 \
--hash=sha256:b3ce300f3644fb06443ee2222c2201dd3a89ea6040541412b8fa189341847218 \
--hash=sha256:b97fe8060236edf3662adfc2c633f56a08ae30560c56310562cb4f95500022d5 \
--hash=sha256:bfe25acf8b437eb2a8b2d49d443800a5f18508cd811fea3181723922a8a82b07 \
--hash=sha256:cd25bcecc4974d09257ffcd1f098ee778f7834c3ad767fe5db785be9a4aa9cb2 \
--hash=sha256:d209d8969599b27ad20994c8e41936ee0964e6da07478d6c35016bc386b66ad4 \
--hash=sha256:d5241e0a80d808d70546c697135da2c613f30e28251ff8307eb72ba696945764 \
--hash=sha256:edd8b5fe47dab091176d21bb6de568acdd906d1887a4584a15a9a96a1dca06ef \
--hash=sha256:f870204a840a60da0b12273ef34f7051e98c3b5961b61b0c2c1be6dfd64fbcd3 \
--hash=sha256:ffa75af20b44f8dba823498024771d5ac50620e6915abac414251bd971b4529f
numpy==2.0.0 \
--hash=sha256:04494f6ec467ccb5369d1808570ae55f6ed9b5809d7f035059000a37b8d7e86f \
--hash=sha256:0a43f0974d501842866cc83471bdb0116ba0dffdbaac33ec05e6afed5b615238 \
--hash=sha256:0e50842b2295ba8414c8c1d9d957083d5dfe9e16828b37de883f51fc53c4016f \
--hash=sha256:0ec84b9ba0654f3b962802edc91424331f423dcf5d5f926676e0150789cb3d95 \
--hash=sha256:17067d097ed036636fa79f6a869ac26df7db1ba22039d962422506640314933a \
--hash=sha256:1cde1753efe513705a0c6d28f5884e22bdc30438bf0085c5c486cdaff40cd67a \
--hash=sha256:1e72728e7501a450288fc8e1f9ebc73d90cfd4671ebbd631f3e7857c39bd16f2 \
--hash=sha256:2635dbd200c2d6faf2ef9a0d04f0ecc6b13b3cad54f7c67c61155138835515d2 \
--hash=sha256:2ce46fd0b8a0c947ae047d222f7136fc4d55538741373107574271bc00e20e8f \
--hash=sha256:34f003cb88b1ba38cb9a9a4a3161c1604973d7f9d5552c38bc2f04f829536609 \
--hash=sha256:354f373279768fa5a584bac997de6a6c9bc535c482592d7a813bb0c09be6c76f \
--hash=sha256:38ecb5b0582cd125f67a629072fed6f83562d9dd04d7e03256c9829bdec027ad \
--hash=sha256:3e8e01233d57639b2e30966c63d36fcea099d17c53bf424d77f088b0f4babd86 \
--hash=sha256:3f6bed7f840d44c08ebdb73b1825282b801799e325bcbdfa6bc5c370e5aecc65 \
--hash=sha256:4554eb96f0fd263041baf16cf0881b3f5dafae7a59b1049acb9540c4d57bc8cb \
--hash=sha256:46e161722e0f619749d1cd892167039015b2c2817296104487cd03ed4a955995 \
--hash=sha256:49d9f7d256fbc804391a7f72d4a617302b1afac1112fac19b6c6cec63fe7fe8a \
--hash=sha256:4d2f62e55a4cd9c58c1d9a1c9edaedcd857a73cb6fda875bf79093f9d9086f85 \
--hash=sha256:5f64641b42b2429f56ee08b4f427a4d2daf916ec59686061de751a55aafa22e4 \
--hash=sha256:63b92c512d9dbcc37f9d81b123dec99fdb318ba38c8059afc78086fe73820275 \
--hash=sha256:6d7696c615765091cc5093f76fd1fa069870304beaccfd58b5dcc69e55ef49c1 \
--hash=sha256:79e843d186c8fb1b102bef3e2bc35ef81160ffef3194646a7fdd6a73c6b97196 \
--hash=sha256:821eedb7165ead9eebdb569986968b541f9908979c2da8a4967ecac4439bae3d \
--hash=sha256:84554fc53daa8f6abf8e8a66e076aff6ece62de68523d9f665f32d2fc50fd66e \
--hash=sha256:8d83bb187fb647643bd56e1ae43f273c7f4dbcdf94550d7938cfc32566756514 \
--hash=sha256:903703372d46bce88b6920a0cd86c3ad82dae2dbef157b5fc01b70ea1cfc430f \
--hash=sha256:9416a5c2e92ace094e9f0082c5fd473502c91651fb896bc17690d6fc475128d6 \
--hash=sha256:9a1712c015831da583b21c5bfe15e8684137097969c6d22e8316ba66b5baabe4 \
--hash=sha256:9c27f0946a3536403efb0e1c28def1ae6730a72cd0d5878db38824855e3afc44 \
--hash=sha256:a356364941fb0593bb899a1076b92dfa2029f6f5b8ba88a14fd0984aaf76d0df \
--hash=sha256:a7039a136017eaa92c1848152827e1424701532ca8e8967fe480fe1569dae581 \
--hash=sha256:acd3a644e4807e73b4e1867b769fbf1ce8c5d80e7caaef0d90dcdc640dfc9787 \
--hash=sha256:ad0c86f3455fbd0de6c31a3056eb822fc939f81b1618f10ff3406971893b62a5 \
--hash=sha256:b4c76e3d4c56f145d41b7b6751255feefae92edbc9a61e1758a98204200f30fc \
--hash=sha256:b6f6a8f45d0313db07d6d1d37bd0b112f887e1369758a5419c0370ba915b3871 \
--hash=sha256:c5a59996dc61835133b56a32ebe4ef3740ea5bc19b3983ac60cc32be5a665d54 \
--hash=sha256:c73aafd1afca80afecb22718f8700b40ac7cab927b8abab3c3e337d70e10e5a2 \
--hash=sha256:cee6cc0584f71adefe2c908856ccc98702baf95ff80092e4ca46061538a2ba98 \
--hash=sha256:cef04d068f5fb0518a77857953193b6bb94809a806bd0a14983a8f12ada060c9 \
--hash=sha256:cf5d1c9e6837f8af9f92b6bd3e86d513cdc11f60fd62185cc49ec7d1aba34864 \
--hash=sha256:e61155fae27570692ad1d327e81c6cf27d535a5d7ef97648a17d922224b216de \
--hash=sha256:e7f387600d424f91576af20518334df3d97bc76a300a755f9a8d6e4f5cadd289 \
--hash=sha256:ed08d2703b5972ec736451b818c2eb9da80d66c3e84aed1deeb0c345fefe461b \
--hash=sha256:fbd6acc766814ea6443628f4e6751d0da6593dae29c08c0b2606164db026970c \
--hash=sha256:feff59f27338135776f6d4e2ec7aeeac5d5f7a08a83e80869121ef8164b74af9
# via
# -r requirements.txt
# torch-mlir
packaging==24.0 \
--hash=sha256:2ddfb553fdf02fb784c234c7ba6ccc288296ceabec964ad2eae3777778130bc5 \
--hash=sha256:eb82c5e3e56209074766e6885bb04b8c38a0c015d0a30036ebe7ece34c9989e9
packaging==24.1 \
--hash=sha256:026ed72c8ed3fcce5bf8950572258698927fd1dbda10a5e981cdf0ac37f4f002 \
--hash=sha256:5b8f2217dbdbd2f7f384c41c628544e6d52f2d0f53c6d0c3ea61aa5d1d7ff124
# via torch-mlir
sympy==1.12 \
--hash=sha256:c3588cd4295d0c0f603d0f2ae780587e64e2efeedb3521e46b9bb1d08d184fa5 \
--hash=sha256:ebf595c8dac3e0fdc4152c51878b498396ec7f30e7a914d6071e674d49420fb8
sympy==1.12.1 \
--hash=sha256:2877b03f998cd8c08f07cd0de5b767119cd3ef40d09f41c30d722f6686b0fb88 \
--hash=sha256:9b2cbc7f1a640289430e13d2a56f02f867a1da0190f2f99d8968c2f74da0e515
# via torch
torch==2.4.0.dev20240428+cpu \
--hash=sha256:2528bdfa7dc56208c06d576f6f9ac69a75caf5b6a2c22ae44c82ac001e2e6f3e \
--hash=sha256:fbae9a5092288805b4eb9d8882481343aa2637bc5c6876ae0c50d93275dafe7d
torch==2.4.0.dev20240604+cpu \
--hash=sha256:95dd17654e0f7c82a9ef50dca328222d7b434b7cc18713f0b83f99bc3f14fb2c \
--hash=sha256:cb1f42d3bcf6b6b4a107a0c6d0b6764fa9ef3136b4575b6136647c703b65c725
# via
# -r requirements.txt
# torch-mlir
torch-mlir==20240430.86 \
--hash=sha256:9f7f2c23a5f21c63b50738202e39809bbb11f32e73edbde6c8b20ec4b0af0d8c \
--hash=sha256:b161e02ff11035a8444a165ed1731ac602718d55ca542bc6b3e513182b8a05b4
torch-mlir==20240620.127 \
--hash=sha256:89b04afe4d39b273cefa141d69bcaf72a5517924d7b87799a88b3fd4c7786e7e \
--hash=sha256:8ef2305b6f6846fddeda6fa9878548392eb7146cbba539db6c7da717b0f68c81
# via -r requirements.txt
typing-extensions==4.11.0 \
--hash=sha256:83f085bd5ca59c80295fc2a82ab5dac679cbe02b9f33f7d83af68e241bea51b0 \
--hash=sha256:c1f94d72897edaf4ce775bb7558d5b79d8126906a14ea5ed1635921406c0387a
typing-extensions==4.12.2 \
--hash=sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d \
--hash=sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8
# via torch
2 changes: 1 addition & 1 deletion test/AotCompile/basic_tcp_ops.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ func.func @func_broadcast_add(%arg0: tensor<f32>,
%arg1: tensor<?xf32>) -> tensor<?xf32> {
%c0 = arith.constant 0 : index
%dim = tensor.dim %arg1, %c0 : tensor<?xf32>
%arg0_ex = tensor.expand_shape %arg0 [] : tensor<f32> into tensor<1xf32>
%arg0_ex = tensor.expand_shape %arg0 [] output_shape [1] : tensor<f32> into tensor<1xf32>
%arg0_bcast = tcp.broadcast %arg0_ex, %dim {axes = [0]} : tensor<1xf32>, index -> tensor<?xf32>
%0 = tcp.add %arg0_bcast, %arg1 : tensor<?xf32>, tensor<?xf32> -> tensor<?xf32>
return %0 : tensor<?xf32>
Expand Down
4 changes: 2 additions & 2 deletions test/Conversion/TorchToTcp/data_movement.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,8 @@ func.func @torch.aten.slice.Tensor(%arg0: !torch.vtensor<[1,56,?,?],f32>) -> !to
// CHECK-LABEL: @torch.aten.gather
// CHECK-SAME: %[[ARG0:.+]]: !torch.vtensor<[1,4,3],f32>,
// CHECK-SAME: %[[ARG1:.+]]: !torch.vtensor<[1,4,2],si64>) -> !torch.vtensor<[1,4,2],f32>
// CHECK: %[[V1:.+]] = torch_c.to_builtin_tensor %[[ARG0]] : !torch.vtensor<[1,4,3],f32> -> tensor<1x4x3xf32>
// CHECK: %[[V2:.+]] = torch_c.to_builtin_tensor %[[ARG1]] : !torch.vtensor<[1,4,2],si64> -> tensor<1x4x2xi64>
// CHECK-DAG: %[[V1:.+]] = torch_c.to_builtin_tensor %[[ARG0]] : !torch.vtensor<[1,4,3],f32> -> tensor<1x4x3xf32>
// CHECK-DAG: %[[V2:.+]] = torch_c.to_builtin_tensor %[[ARG1]] : !torch.vtensor<[1,4,2],si64> -> tensor<1x4x2xi64>
// CHECK: %[[GATHER:.+]] = tcp.gather %[[V1]], %[[V2]] {dim = 2 : index} :
// CHECK-SAME: tensor<1x4x3xf32>, tensor<1x4x2xi64> -> tensor<1x4x2xf32>
// CHECK: %[[V3:.+]] = torch_c.from_builtin_tensor %[[GATHER]] : tensor<1x4x2xf32> -> !torch.vtensor<[1,4,2],f32>
Expand Down
Loading
Loading