Programming Model for Export - PyTorch Compiler Series Episode 1

Key Takeaways

All right. Today we are going to talk about the torch export uh programming model. Um so before we go into what a program model is um let's let's give some background on what to touch export is trying to do. Essentially it's trying to take a model and create a graph representation of the model. um and um exactly how it does it and and what kind of guarantees it it's uh it's intending to to get out of this process uh will inform what the program model is right uh like so so let's let's look at like uh the the most general picture of what's happening with the model you have a space of inputs and if you take some inputs um and run them through your model you're supposed to get some outputs right so this is what uh we we're taking taking as a baseline uh which is eager execution of your model. All right. So um so for export what you do is you you take an example input which is supposed to be representative of of of a subset of your of your inputs and you export your model with this example input and uh what you get is a graph. So, so conceptually you can think of this as some subset of your model or some subset of the paths through your model which um this example input um exercises right but again it's supposed to be representative of uh of uh various other inputs that uh that would go down the same path uh through your model. So uh essentially what um so so then the graph that you get should work exactly like your original model on on a large subset of your inputs. Right? So that's that's what the the last picture uh shows is that now you have extracted a subset a subspace of inputs uh kind of that that is determined by the graph itself. And um and if you if you take the same inputs that that you wanted to run your original model through uh you will get the same outputs right. So that that is the high level intent behind export and frankly it should be the uh the model that um any other such graph uh capturing mechanism should should follow. Right? So so what are the basic principles of graph capture? Uh really it's about balancing safety versus coverage. So uh safety essentially boils down to am I behavior uh behaviorally kind of u uh similar to to the original model and coverage means that you you are able to run this process through a large u set of models right so so uh you you don't you don't have to sort of uh be very pedantic about how you write your models in order to go through this process right so so Now one of the key observations is that uh any such graph capturing mechanism ultimately how you design it comes down to an important decision on what can change after your exporting uh versus what is fixed at export time. So um colloally you would say you know what is dynamic and what is static. you you need to sort of this is a design decision that you want need to make when you're when you're designing graph capture mechanism. All right. So um so if if if safety and coverage are are your two axes you can plot how uh various technologies that try to do this that that came before touch export where they lie in the spectrum. Uh so if you look at jet trace it had high coverage by which by which we mean it will successfully produce a graph for for a large set of models but it is low on safety as in uh it doesn't really define what subset of inputs uh would give you valid outputs after after you uh get your graph. Uh so so if if you for example give inputs that did not actually go through the that path um through your original model you might still execute um the graph as if it did and you would get incorrect results. JIT script was in some sense the opposite uh in in the opposite side of the spectrum where uh its coverage is pretty low. you have to work really hard to to sort of be constrained by the the parameters of what just script is happy with when you're writing a model. But once you do that and it passes just script and you get a graph um it's it's likely to be safe, right? So um FX trace is another intermediate thing that that the PyTorch team developed uh which uses symbolic tracing instead of uh sort of tracing with with u exactly the inputs that you give. So it it it generalizes the inputs a bit more. So it it covers uh more models than JIT script because it's easier to use but it also sort of uh is not as safe as uh JITScript. It's it's safer than JIT trace though. Um so where is export on this spectrum? Well, export tries to have a design that uh gives you high high levels of safety as well as high levels of uh coverage. And it comes down to again being very careful about what can change versus not change at export time uh after export time. So it it the the principles are are shown on the slide. Basically the the first principle is that the graph must not depend on tensor data because we assume that the tensor data will change across runs. There's no way that you will run with exactly the same tensors that you pro provided um when uh when you provide these example inputs. Uh on the other hand um we found that it is useful to specialize based on tensor shapes. So tensor shapes would uh would be an abstraction over your tensors that that says okay the data might change but the shapes will not change across runs. So um we can specialize and the more we specialize in some sense it is is correlated with how easy it is to export your program and also you can do optimizations based on these specializations. Another decision that export makes is that anything that is unrelated to tensors uh they can be fixed. So this includes constants in your program uh which are sort of burnt into the graph. Uh this also covers um any any containers that you use um Python level containers that you use in your in your uh in your program uh their structure is fixed. So anything that is not tensor is essentially specialized. All right. So rather than explain to you um what the program model sort of rules are uh in some sense it's easier to to just show you a series of examples that that help you have some mental model of what the program model is like like and the benefit of a program model is that you can use this understanding to predict what would export do uh on your program right so that that ultimately that's that's our goal for this video. All right. So let's look at some examples now. So uh the first example we will look at is uh is this constant specialization idea. Uh so if there are any constants in your program, they will they will be enlined and folded in. Uh and so this example shows you what will happen. Um let's say if if your model has a tensor which as we said it's assumed to be dynamic but you also take an int and here we pass in the constant three uh at export time and um what will happen is that the it it will go through a and then uh b is a + 7. So it will actually get uh evaluated at export time to 10. And if you if you look at the graph that um is produced in this case you would see that um 10 as a constant appears. So that's 3 + 7. All right. So so here we are taking um what does export uh expect? It takes the the model instance and and your inputs and it produces some graph that uh x is kept as x but a uh whatever the value that you gave at export time um it sort of used constant evaluation to to figure out what what should go into the graph. All right. So uh yeah I also talked about um Python level containers that you can use uh and how their structure is also specialized right. So here we have another example of a model that uh takes a list of tensors and a list of ins and basically does a very similar thing as the previous example. um it it sums over the list of constants and then it adds that to the to a sum of tensors um that you also provided as a list. Right? So um at export time you decide to provide a list of uh two tensors and and a list of three constants. Um and if you look at the graph that was produced in this case um you will see that um it okay so so you will see that it created a graph with five inputs. Two of them are tensors and three of them are constants. Um and it did the sum over the list basically the sum over the two tensors and it added the list of constants but it actually constant evaluated the the constants right so so you provided 3 4 5 they got summed to 12 and you see the see the number 12 in the graph right so this gives you an idea of what um at export time gets gets burned into the graph and what doesn't right so the structure of containers gets burnt in constants get burned and the tensors themselves do not get burnt. All right. So, uh this is a good place to maybe maybe um sort of have a digression and say that um okay, maybe I don't want to use just vanilla primitive uh built-in types. I I want to have a data class that represents my inputs. And so this this example is exactly like the previous example but I am packaging up my inputs in a data class and my model takes data uh in the form of that data class and then it does exactly what it what we did before. So what happens here? Uh well uh if you just try to run this um as is this will not uh succeed export because export does not know what what your data class is and in particular it cannot flatten it into that kind of list of inputs that that we saw. Right? In the previous example we had five inputs. So we knew how to kind of sort of flatten the lists out into into individual things here. We don't uh don't know that. But um it's a oneline change to your to your model or to your exporting workflow to to do this. So use um a registration mechanism that we provide for data classes. Uh there's a more generic registration mechanism for arbitrary classes as well. Um so if you if you do this then we know how to inspect into a data class and sort of flatten it out. So so now you see that you get exactly the same graph as in the previous example where you have five inputs two of them are tensors and so on. So so it just basically opened up the data class and and did it thing and it still specialized u it still evaluated the constants. Okay. So what else is um is specialized? So let's let's look at that. So uh I mentioned in the in the principles that uh while we do not specialize on tensor data, we do specialize on tens tensor shapes. So in this case um what you're doing is instead of taking a list of constants you're taking a list of tensors but we are extracting the the shapes out of those tensors and doing some u evaluation on those shapes u in particular we're taking the sum of the shapes and then we're doing exactly what we were doing before. So here instead of passing 345 you're passing in tensors of shapes 345. And what will happen in this case is uh I mean should be obvious if if you're following what I'm saying which is the shapes would be treated as constants and you would get uh get the same graph that you were getting in previous examples. Right? Uh also I printed out uh the graph that we produce um in some more detail. Um so if you use uh EP um the exported program has a graph module inside. So if you if you use print readable on it, you will get some more information what the what the shapes the expected shapes of the inputs are. So in this case we we used shapes two and two for the the list of tensors and then for the constants that were for the for the shapes that for for the list of tensors for which we were extracting the shapes. You can see that we are expecting tensors of shapes 3 4 5 here. Um so you can get that extra information. And you can also get some metadata on on which line user code line produced some snippet of code. All right. So um so that's that's what we do for shapes. We we specialize them. Um but we also do some more on shapes. So so just like you can you can do manipulations and constants and we would follow along. there is a mechanism to to propagate shapes through through your other operators as well. So here we are using um torch.cat instead of summing the shapes we are concatenating the given list of tensors and taking their shapes. So um so if you if you if you see what what happens here you will see that again we we found out that the sum of the shapes was 12 but now we are getting to the sum in an indirect way. So how are we doing that? So torch.cat would basically concatenate the given tensor. So, so if you pass in shapes 3 4 5 it will get a new tensor of shape 12 and you can see that as an intermediate like the shape of an intermediate value which is the output of the the torch cat operator and then the rest of it follows along. So um so this is this is very very nice because you know that inside the PyTorch compiler is actually tracking the shapes of individual operators not just the inputs but but it can uh sort of continue the reasoning through your model. All right. So next example is uh something else you can do which is also quite common uh which is you can have control flow uh on shapes and because we are uh we understand shapes at compile time we can actually understand which branch to to go to go and and produce continue to produce a graph right so um export is a tracing based uh compiler so we are actually going through certain parts in your program as we and and we are capturing them. Um so so in this particular example you passed in shapes 3 4 5 and we concatenated them to to a tensor of uh of shape um oh so sorry so so here I was going to hide these inputs. So so let me run this program again. Okay. So, initially we are passing in shapes 3 4 5. Yeah. So, they are uh the concatenation of them will have shape uh 12. So, it will actually go down the else branch in this case and trace this uh and you can see that happening here. Right? So, it basically output the code for the else branch. Uh this is this is the code that that that is going that is equivalent to that that line line 15. And um if you want your uh graph to capture a different path you could give different inputs. So for example here you're giving shapes 6 7 8 and now uh the concatenation of these tensors will have shape 21. And so it will uh go through the true branch now because the the condition was that uh it should be um the the shape should be more than 20 and in this case it's 21. So yeah, if you if you re rer rerun the export now you're seeing that the concatenation indeed has 21 and so now it has traced the true branch which is uh this line right and so you get an idea that the graph would be specialized to the shapes that you provide at uh at input time. So, so your again your example input is supposed to be representative of the set of inputs you care about uh for inference and so depending on what those what that subset is you you might get different graphs. Um and so this this kind of sets expectations on that. All right. So so this immediately makes you think of like okay what if I want to capture the graph in a more generic way and not really specialize that much on shapes. And this is where dynamic shapes come in. So to illustrate that, let's look at let's pop back and like remove the complexity from from the from the model and and think of a very generic kind of uh operation like add mm. Right? So you you take three tensors xyz and I'm going to um matrix multiply the y and the z and and add x to it. So that's that's what this this forward method does, right? And um I'm going to give example inputs that fit this. So so ultimately I'm going to multiply a 3 + 4 and a 4 + 5 getting a 3 + 5 and so that when I add it to add the result of the matrix multiply to to x I'm going to get correct shapes. Right? So, so uh you could export this program without using dynamic shapes of course and then then you would the graph would be specialized to the shapes three and four and five but here we want to do something more generic. Um so um how do you specify dynamic shapes to export? Uh well there are many different ways. One convenient way is that you can create a shapes collection uh and just add shapes of your input tensors to that. So, so you're going to use XYZ as your example inputs. You can specify that each of them are two-dimensional uh tensors that have uh dynamic shapes. And the way you do that is by saying the they're dimnamic. Um so if you if you try to um export this program um the the exported like the um oh so I'm sorry. So this is another example where I want to show more things later. So let's uh let's not get ahead of ourselves. So let's let's just take uh try to export the thing that we we said. Okay. And then um in in the input shapes you will see some symbols uh coming up now. So instead of 3 4 5 you're getting these symbols that represent the generality of this of this model. Right? So let's say y has shape s0 cross s3 and z has s3 cross s1. Then the result will be of type s0 cross s1. And then you can add it to another s0 cross s1 which which x is now constrained to to be right. So so the mat mole as I said will be as 0 + cross s1 and then uh x is of the right shape. So so the addition will also be as zero cross s1 and everything succeeds. Okay. So, so now if you if you if you think about what happened behind the scenes, right? So, we while we are tracing we as I said we are propagating shapes and now but but now we are propagating shapes in in the symbolic form and we know that uh certain shapes have to be equal uh for some operations to succeed. So uh the symbolic reasoning that is going behind the scenes is kind of constraining the system so that so you get the right right shapes out. So this is kind of like type checking uh but except that it's happening at the level of tensor shapes. All right. So um so yeah the initially the errors that we are seeing was because I wanted to show some more variations of this program. So, so one of the things uh you can use to test whether the the graph is truly truly kind of generic in dynamic shapes or whether it was overly specialized to the to the shapes that you provided which was at the end of the day is still concrete shapes like 3 4 5 right. So you can you can try to do this uh run this um program with different shapes right so you can say that okay we are going to everything is going to be six cross uh six tensors now Z would be some random tensor and we going to multiply it with entity matrix and so we are going to get back the same tensor and if you add the negative of the tensor uh we are expecting all zeros so so if you run this program it should still run fine. Uh the assertion should uh should succeed. On the other hand, if you changed one of the inputs to be of different shape but you did not change the other inputs, right? So now you have sort of shape mismatch. Uh the entity matrix is now 7 + 7. Uh this shouldn't work and you would expect that this will obviously fail at runtime. But uh what you get is yeah also a runtime failure but it's a very early um error. So even before your your model was run, your explorer program was run um we have some checks going on that um you know ultimately the the subset of inputs that you had in mind they were constrained to be of those dynamic shapes and in this case you just cannot find good assignments of S0, S1 and S3 that work and so so that is what causes this equivalent of a um sort of like a runtime type error uh which is that we expected some shape to be six but we got seven right uh what is args one args one well the graph has forgotten the names of the original kind of inputs so args one is is basically y um and y dotshape uh is expect to be six because x was given to be six and um and z is still six. So the seven is the is the one that is a problem here and you get uh get an error message very early on uh when when you're trying to execute this graph. All right. So uh dynamic shapes also get propagated. This is shouldn't be a surprise at this point but we just want to show that as an example. uh instead of passing in one Z if you if you passed in sort of a list of tensors uh and you concatenated them and continued with that mm operation you can uh of course the example inputs have to now uh be arranged such that you know you don't get any obvious mistakes. So here we are still expecting the uh the concatenation to to have shape 4 + 5. So one way you can you can do that is by having two tensors um um each each of shape 2 + 5 right. So um so this should work um you specify the dynamic shapes uh for all the tensors in your inputs. Everything let's say is dynamic to begin with. Um here you're not really saying which shapes should be equal. You're letting the the system figure it out. Um, and so this this should work. Um, so let's see what happens here. All right. So just as before we we kind of flattened out the list of tensors to two inputs here. The shapes are a little bit more complicated now. So, so you have in general two tensors S4 cross S1 and S6 cross S1. Um, and the concatenation has a shape S4 plus S6, comma S1, right? That's what you expect your Y to satisfy in the in um in the second dimension. Uh and then yeah so so all of the the shapes getting propagated you can inspect the intermediate nodes in your graph and and you can check that this works and again everything is dynamic here. So if you if you constructed some other inputs as long as they satisfy this this pattern or shapes that that you you expected in your in your um in your in your generated graph it should be fine. So, so here we are concatenating two tensors uh whose first dimensions are 2 and four. So, they will produce a 6 + 6 matrix and everything will be fine. So, so this means that the assertion should also succeed. And actually, yeah, I I didn't I I I ran the assertion as well as part of uh the initial run. So, so the fact that it passed tells you that yeah, it worked. Okay. What else can happen uh here? So, so this is a more even more interesting example. Um, here we have more or less the same program. Uh, but now we are sort of taking a slice of one of the input tensors. Uh, we're taking the first three um, you know, we are we are slicing the the original X tensor and and taking the first three um rows. So um so what do we what do we expect to happen here? Um let's see. I don't want to run things too much. Uh I just want to do the export. So um so we pass in a tensor that has a lot of rows but we know that only the first three will be relevant. And uh and if you choose to add that to the rest of the uh matrix multiplication, we we expect Y to also have the first dimension be three. So that's why we're doing we're passing in this this value uh this this this shape of tensor and and ZS again will be concatenated. So it should match up with the six. So 2 + 4 is six. So this this should work. We made one small change in our dynamic shape spec. um y has no choice other than being uh being three here because of this hard coding of the index value three. Right? So if you passed in dynamic, we would actually error here because if you if you expect a dynamic shape and it turns out to be like specialized to a constant, then we would say okay you you expected your uh shape to be dynamic but it's not right. It actually got specialized to the value two. So I'm not going to show you the error that that happens here, but the the suggested fix that will pop out is that you need to make the static, right? You or or you can pass in literally three as a shape here. So I made that change to make the export go through and uh you'll see that it it works. Now I'm also printing out a bunch of other information that I'll explain in a moment. So yeah the the the shape the names changed um but the same thing is happening right so you have S3 cross S1 and S5 cross S1 and you have S3 plus S5 here showing up um there's an initial slice operation and you can see that the first dimension has to be three and that that shape is actually specialized and burned into the graph and uh one other way you can you can get to this information is that here is a little snippet of code I wrote where you are going through the nodes in your graph and matching on nodes that have name x. Here I know that only the input node will have have name x. Uh each node carries some metadata that will give you the shape and this is how we actually print out these shapes uh when we are printing out the graph. And then there's another part of the export program that records some range constraints. So uh just pop back here and think about this right. So we passed in a 7 + 5 and we were successfully able to index into into that. Uh what if we pass something uh where whose first dimension was less than three then the index operation would not succeed. So actually when we are exporting we are not only propagating the shapes but we are keeping track of the bounds on those shapes as well. So if I print out what is the lower bound of this symbol as zero we will see that it is at minimum it has to be three. So that that makes sense right because um no matter what other values you you provide to x the first dimension has to be at least three for for the for the dimensions uh for for the slice to to work. So our example inputs match that. Um if we if we try to run this uh code then um we would get an error because we said that the first dimension of y has to be three. So that's one kind of error we are used to where we burnt in the value three into the graph and you provided a six. So so that's not good. So we will get an early error on that. Um likewise if you passed in an x where the first dimension was too little it's it's it's not uh a minimum of three then the early check will also tell you that right so here we are saying that okay I expected ar0 ar0 is x that shape to be greater than equal to three but we got two right so you get these kinds of early checks and these are a result of the shapes shape propagation that we All right, another thing uh to keep in mind is that your model can have uh state and state will all always be assumed to have static shapes. Um there's no way you can make them dynamic. Uh and if they interact with the rest of your model, then they can force other shapes to become static as well. uh this this makes sense because your your model can sort of mutate state and do other funny things with it and um and it's impossible to to track all of that through one run of your program necessarily. So, so we do not allow your states to have dynamic shapes. And so the effect of this on this particular program is that um you will yeah you will you will get some static shapes. So so your state had shape 4 + 4. So you will see that 4 + 4 appear somewhere. Um but note that some other shapes were are still dynamic. Um what we did here was we instead of instead of deciding which shapes should be dynamic or static, we use some other thing called dim.auto which which allows your uh shapes to be either static or dynamic, right? And it just figures out which one is which by itself. So that's that's a bit more convenient. Uh of course, if you want more control, there are other ways of specifying dynamic shapes where you can name them and give them bounds and so on. I'm just showing you some kind of very very quick ways of um of doing this while letting the system figure out most of it. So the fact that we provided some sixes here as first dimension that was not so important to this program. So that can that can still be dynamic, right? And so in principle you can you can pass in you you can change these shapes to be 7 8 5 whatever you want but the second dimension has to be four to line up with the with the state right that's what you would expect. All right so let's switch gears to something else that is um comes back to this idea that your um your graph should not depend on data because data is expected to change over time. Right? So, so this is uh an example of a model where uh you you take in two tensors but now you are you're doing this. Okay, first you're adding the the two tensors. This is just a device I'm using to uh to basically make the shapes of X and Y the same. Um and you'll you'll see why in a moment. So I want to kind of make a point and and so so this is a convenient way of doing it but maybe your program actually did it. Anyway, the important bit is here that you're using an operation let's say count nonzero that uh what it does is it counts the number of non-zero items in your um in your input tensor X and this is inherently a data dependent operation. So if you passed in another tensor be it of the same shape but it could have a different number of zero items right. So so the count non-zero is a data dependent operation uh it returns to you a zero dimensional tensor. So you have to do an item to actually get a int out and then we are doing the indexing on using that um int. Right? So you're basically extracting the first C values out of C rows out of Z the the sum of X and Y. Um for exporting I' I've stripped out all the dynamic shapes here. So um so you you are passing in two uh tensors each of you know there are six rows in it or six values in it. And so let's see what happens if you're trying to export this. This is not this is not um we don't expect this to work because we have no idea what is going to happen when we are doing this indexing right. Um actually we do but but that that we'll we'll get to that a little bit later. So what's happening here is that we are getting an error which is the the much fear data dependent error that we get from export and it is saying something about like some some expression could not be resolved like we don't know what what whether it's less than zero or not and what this u0 is um it's a different kind of symbol than the szer's and the s1's of the uh that we uh saw in previous examples Because anytime you have a U that means that something data dependent happened and we cannot really track it or propagate it during export because we we erase the data when we are exporting. So so that we are never dependent on it. But there are some suggested fixes. So um it says the error is in the indexing. So now you can understand why it was worried whether the index could be negative or or non- negative. because uh PyTorch actually allows you to negative index um tensors and so but but it has no idea of what what can happen here. So it does suggest that you you should you should either assume that C is non- negative or it's negative. We know in this case that it must be non- negative, right? Because it's counting the number of zeros in your tensor. So, so we can add that uh something called torch check which basically adds an assumption to um to to PyTorch and then it turns into a runtime assertion. Um so if you if you run your program if you run your export again now you get a different error. Now it's it's still related to the indexing but now it's saying is it less than equal to six? Is it is it less than the the bounds of of Z because Z was inferred to be a 6 cross uh 6 + one as well, right? Um because your example inputs said that. So it's another tort check that you have to add and uh and now your program should work. So um so now you get a successful export. the values six have been burned in uh because we did not use dynamic shapes. Uh using dynamic shapes would be orthogonal to this example. The important bit is that um the result of nonzero right it's a zero dimensional tensor. If you take the item uh if you take an item out of it we do not know what what value it will have. So it is a symbolic int. That's what sim means. And it's named by by this symbol u0. We do not know anything apirally about about this symbol whether it's negative positive less than what or whatever. So but because we added a bunch of checks there are two runtime assertions that we added that it needs to be non- negative and it it is less than equal to six and then we we do the slice right. So, so this is what happened here. But, but now zoom out a little bit. Now, you're not happy with with this situation. Why? Because you had to add a bunch of tort checks. Well, it would be okay if you needed to add the torch checks if they were really justified. If if really the pyarch compiler could not figure out ever what was going on. But in this particular case uh you know you know that your mental model says that you should know that X has uh shape six and if you are counting the number of non-zero values in X it better be between 0 and six. So why do we need to add these checks right? Well, it has to do with um what we specify um the behavior of count non-zero is during during export. So that is driven by these so-called fake ops which for every actual op we have another implementation of the op uh which tells you how to calculate the output shapes given the inputs. So we used to doc cat before. So it knows that the output shape has to be the sum of the the input shapes. And that's how you got the the I don't know S3 plus S5 or whatever in a previous example for count non zero. It says well the output is I don't know what it is right and actually you're doing an item on it and and it has lost all track of what the input tensor was at that point. So we can we can potentially do something useful here. we can introduce something called custom ops to fix this problem. So what we're going to do here is write our own version of count nonzero and we are going to make it clever and one way to do it is is is to write a custom op and so we are writing a custom op here which is let's say we call it count nonzero int and um and we we are going to redirect to the original count nonzero here we are also going to do a dot item on it so that the result is actually an int zero dimensional tensor on which you have to do a dot item. But true to form whenever you have a op uh you have to register a fake implementation of that that is actually used during export. So remember that export actually it doesn't run your program with uh with real tensors. it it erases the data part of the tensor so that it's not dependent on the data. So but but now it it cannot actually run count non-zero on it. It will run this fake implementation to to realize what the shape of the output should be so that it can continue to trace. So here we are going to say that okay the output of this count non-zero int is going to be some int right. So, so the way you create an unknown int is by using new dynamic size from uh from PyTorch. And now we have special knowledge that the output has to be between zero and the number of elements of the input. Um right so so we can ahead of time add these dot checks to satisfy the compiler. So this this can be done once and now in your in your actual program you are free to slice with this um with this new operation right and so what what will happen here is that it will know that the number of non-zero items in X has to be between 0 and six right um but but the Z uh tensor a period it's unrelated to X And this is why we we actually added X and Y so that we can infer that Z has the same shape as X. Um so really we in some sense we have given the PyTorch compiler all the information it needs that that this slicing should be safe because Z will have the same shape as X. We know that the output of count non-zero int must be between 0 and six which is the shape of X. So this should not produce an error, right? And so that's that's what should happen. So let's see if this happens. Okay. So so yeah. So so this this passes um and in fact now we can go ahead and say there was nothing inherently special about the value six here. uh note that in the implementation we used x dot number of elements. So we we can generalize this to whatever we want right. So in particular we can say that um x and y can be of dynamic shape um and this code should still work right. So in this particular iteration uh we specialized the value six. But now if we rerun this program and I uncmp uncommented the the dynamic shape spec we should get a program where it's more generic. So so in this case we we will get um another explorer program and now uh the input shapes can be dynamic right. So uh yeah that's that's all I had for today. Uh so hopefully this gives you a uh good idea of um what export program model is like. It basically boils down to very simple choices uh at the design level of what is dynamic and what is static and much of the rest of the behavior of export follows from there. Um so yeah that's that's a wrap for today. Hope you enjoyed the video. Thank you.